typos

jaxmarl/environments/jaxnav/README.md

@@ -9,10 +9,10 @@ JaxNav was first introduced in ["No Regrets: Investigating and Improving Regret
 ### Map Types
 The default map is square robots of width 0.5m moving within a world with grid based obstacled, with cells of size 1m x 1m. Map cell size can be varied to produce obstacles of higher fidelty or robot strucutre can be changed into any polygon or a circle.
 
-We also include a map which uses polygon obstacles, but note we have not used this code is a while so there may well be issues with it.
+We also include a map which uses polygon obstacles, but note we have not used this code in a while so there may well be issues with it.
 
 ### Observation space
-By default, each robot recieves 200 range readings from a 360-degree arc centered on their forward axis. These range readings have a max range of 6m but no minimum range and are discritised with a resultion of 0.05 m. Alongside these range readings, each robot recieves their current linear and angular velocities along with the direction to their goal. Their goal direction is given by a vector in polar form where the distance is either the max lidar range if the goal is beyond their "line of sight" or the actual distance if the goal is within their lidar range. There is no communication between agents.
+By default, each robot receives 200 range readings from a 360-degree arc centered on their forward axis. These range readings have a max range of 6m but no minimum range and are discretised with a resultion of 0.05 m. Alongside these range readings, each robot receives their current linear and angular velocities along with the direction to their goal. Their goal direction is given by a vector in polar form where the distance is either the max lidar range if the goal is beyond their "line of sight" or the actual distance if the goal is within their lidar range. There is no communication between agents.
 
 ### Action Space
 The environments default action space is a 2D continuous action, where the first dimension is the desired linear velocity and the second the desired angular velocity. Discrete actions are also supported, where the possible combination of linear and angular velocities are discretised into 15 options.
@@ -52,6 +52,8 @@ viz.animate("test.gif")
 
 ## TODOs:
 - remove `self.rad` dependence for non circular agents
+- more unit tests
+- add tests for non-square agents
 
 ## Citation
 JaxNav was introduced by the following paper, if you use JaxNav in your work please cite it as:

jaxmarl/environments/jaxnav/jaxnav_env.py

@@ -229,7 +229,6 @@ def step_env(
         old_goal_reached = agent_states.goal_reached
         old_move_term = agent_states.move_term
         map_collisions = jax.vmap(self._map_obj.check_agent_map_collision, in_axes=(0, 0, None))(new_pos, new_theta, agent_states.map_data)*(1-agent_states.done).astype(bool)
-        # map_collisions = self._check_map_collisions(new_pos, new_theta, agent_states.map_data)*(1-agent_states.done).astype(bool)
         agent_collisions = self.map_obj.check_all_agent_agent_collisions(new_pos, new_theta)*(1- agent_states.done).astype(bool)
         collisions = map_collisions | agent_collisions
         goal_reached = (self._check_goal_reached(new_pos, agent_states.goal)*(1-agent_states.done)).astype(bool)

jaxmarl/environments/jaxnav/maps/grid_map.py

@@ -346,6 +346,11 @@ def scale_coords(self, x):
     def dikstra_path(self, map_data, pos1, pos2):
         """ Computes shorted path (if possible) between `pos1` and `pos2` on the grid specified by `map_data`.
         
+        Method TL;DR: Dijkstra's algorithm. JIT'd while loop through the open set of nodes, updating the distance
+                      to go for each neighbour in `_body`. Terminates when the open set is empty.
+        
+        TODO: also terminate if the goal is reached. Add unit tests.
+        
         Output:
         - `valid` (bool): True if a path exists
         - `d` (float): distance of shortest path, INF if no path exists
@@ -492,6 +497,8 @@ def plot_agent_path(self,
 ])
 
 class GridMapPolygonAgents(GridMapCircleAgents):
+    """ Map for homogenous, convex polygon agents. 
+    """
 
     def __init__(self,
                  num_agents: int,
@@ -513,14 +520,9 @@ def __init__(self,
         side_length = jnp.maximum(max_x - min_x, max_y - min_y)
 
         cell_ratio = side_length / self.cell_size
-        # print('cell ratio', cell_ratio)
 
         self.agent_window = jnp.ceil(cell_ratio*2).astype(int)  # NOTE times 2 should be enough
-        self.idx_offsets = jnp.arange(-self.agent_window, self.agent_window+1)
-        # print('side length', side_length)
-        # print('agent window', self.agent_window)
-        # print('idx offsets', self.idx_offsets)
-
+        self.idx_offsets = jnp.arange(-self.agent_window, self.agent_window+1)        
 
         #  2D with one set of coords for all agents 
         assert (len(self.agent_coords.shape) == 2)
@@ -604,7 +606,6 @@ def _checkGrid(self, x1y1, x2y2, grid_idx, map_grid):
 
         def _checkLineLine(x1, y1, x2, y2, x3, y3, x4, y4):
             """ Check collision between line (x1, y1) -- (x2, y2) and line (x3, y3) -- (x4, y4) """
-            # TODO understand this
             uA = ((x4-x3)*(y1-y3) - (y4-y3)*(x1-x3)) / ((y4-y3)*(x2-x1) - (x4-x3)*(y2-y1))
             uB = ((x2-x1)*(y1-y3) - (y2-y1)*(x1-x3)) / ((y4-y3)*(x2-x1) - (x4-x3)*(y2-y1))
             c = (uA >= 0) & (uA <= 1) & (uB >= 0) & (uB <= 1)
@@ -653,7 +654,13 @@ def _checkPointRect(x, y, rx, ry):
 
     @partial(jax.jit, static_argnums=[0])
     def check_all_agent_agent_collisions(self, agent_positions: chex.Array, agent_theta: chex.Array, agent_coords=None) -> chex.Array:
-        """ Using Separating Axis Theorem (SAT) to check for collisions between convex polygon agents. """    
+        """ Use Separating Axis Theorem (SAT) to check for collisions between convex polygon agents. 
+        
+        Separating Axis Theorem TL;DR: Searches for a line that separates two convex polygons. If no line is found, the polygons are colliding.
+                                       To search for a separating line, we look at the normals of the edges of the polygons and project the vertices
+                                       of all polygons onto these normals. If the projections of the vertices of one polygon do not overlap with the
+                                       projections of the vertices of the other polygon, then the polygons are not colliding.
+        """    
 
         def _orthogonal(v):
             return jnp.array([v[1], -v[0]])
@@ -696,7 +703,6 @@ def check_agent_beam_intersect(self, beam, pos, theta, range_resolution, agent_c
         @partial(jax.vmap, in_axes=(None, None, 0, 0))
         def _checkSide(beam_start, beam_end, side_start, side_end):
             """ Check collision between line (x1, y1) -- (x2, y2) and line (x3, y3) -- (x4, y4) """
-            # TODO understand this
             x1, y1 = beam_start
             x2, y2 = beam_end
             x3, y3 = side_start