viz example

Author: amacrutherford

jaxmarl/environments/jaxnav/README.md

@@ -1,9 +1,11 @@
 # 🧭 JaxNav 
 
-2D geometric navigation for differential drive robots. Using distances readings to nearby obstacles (mimicing LiDAR readings), the direction to their goal and their current velocity, robots must navigate to their goal without colliding with obstacles.
+2D geometric navigation for differential drive robots. Using distances readings to nearby obstacles (mimicing LiDAR readings), the direction to their goal and their current velocity, robots must navigate to their goal without colliding with obstacles. 
 
 ## Environment Details
 
+JaxNav was first introduced in "No Regrets: Investigating and Improving Regret Approximations for Curriculum Discovery" with an in-detail specification given in the Appendix.
+
 ### 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.
 
@@ -19,13 +21,41 @@ The environments default action space is a 2D continuous action, where the first
 By default, the reward function contains a sparse outcome based reward alongside a dense shaping term.
 
 ## Visulisation
+Visualiser contained within `jaxnav_viz.py`, with an example below:
+
+```python
+from jaxmarl.environments.jaxnav.jaxnav_env import JaxNav
+from jaxmarl.environments.jaxnav.jaxnav_viz import JaxNavVisualizer
+import jax 
+
+env = JaxNav(num_agents=4)
+
+rng = jax.random.PRNGKey(0)
+rng, _rng = jax.random.split(rng)
+
+obs, env_state = env.reset(_rng)
+
+obs_list = [obs]
+env_state_list = [env_state]
+
+for _ in range(10):
+    rng, act_rng, step_rng = jax.random.split(rng, 3)
+    act_rngs = jax.random.split(act_rng, env.num_agents)
+    actions = {a: env.action_space(a).sample(act_rngs[i]) for i, a in enumerate(env.action_spaces.keys())}
+    obs, env_state, _, _, _ = env.step(step_rng, env_state, actions)
+    obs_list.append(obs)
+    env_state_list.append(env_state)
+    
+viz = JaxNavVisualizer(env, obs_list, env_state_list)
+viz.animate("test.gif")
+```
 
 ## TODOs:
-- remove self.rad dependence
+- remove `self.rad` dependence for non circular agents
 
 ## Citation
 JaxNav was introduced by the following paper, if you use JaxNav in your work please cite it as:
 
-'''bibtex
+```bibtex
 TODO
-'''
+```

jaxmarl/environments/jaxnav/jaxnav_env.py

@@ -426,15 +426,7 @@ def _check_map_collisions(self, pos: chex.Array, theta: chex.Array, map_data: ch
     @partial(vmap, in_axes=(None, 0, 0))
     def _check_goal_reached(self, pos: chex.Array, goal_pos: chex.Array) -> bool:
         return jnp.sqrt(jnp.sum((pos - goal_pos)**2)) <= self.goal_radius 
-        
-    @partial(vmap, in_axes=(None, 0, None, None))
-    def _check_agent_collisions(self, agent_idx: int, agent_positions: chex.Array, dones: chex.Array) -> bool:
-        # TODO this function is a little clunky FIX 
-        z = jnp.zeros(agent_positions.shape)
-        z = z.at[agent_idx,:].set(jnp.ones(2)*self.rad*2.1)  
-        x = agent_positions + z
-        return jnp.any(jnp.sqrt(jnp.sum((x - agent_positions[agent_idx,:])**2, axis=1)) <= self.rad*2) 
-    
+            
     @partial(jax.jit, static_argnums=[0])
     def get_obs(self, state: State) -> chex.Array:
         obs_batch = self._get_obs(state)

jaxmarl/environments/jaxnav/jaxnav_viz.py

@@ -4,14 +4,14 @@
 import matplotlib.animation as animation
 from typing import Optional, List
 
-from .jaxnav_env import JaxNav
+from .jaxnav_env import JaxNav, State
 import jax.numpy as jnp
 
 class JaxNavVisualizer(object):
     def __init__(self, 
                  env: JaxNav, 
                  obs_seq: List, 
-                 state_seq: List,
+                 state_seq: List[State],
                  reward_seq: List=None,
                  done_frames=None,
                  title_text: str=None,
@@ -77,8 +77,9 @@ def update(self, frame):
         if self.plot_path:
             for a in range(self.env.num_agents):
                 plot_frame = frame
-                if self.done_frames[a] < frame:
-                    plot_frame = self.done_frames[a]
+                if self.done_frames is not None:
+                    if (self.done_frames[a] < frame):
+                        plot_frame = self.done_frames[a]
                 self.env.map_obj.plot_agent_path(self.ax, self.path_seq[:plot_frame, a, 0], self.path_seq[:plot_frame, a, 1])
             # self.ax.plot(self.path_seq[:frame, 0], self.path_seq[:frame, 1], color='b', linewidth=2.0, zorder=1)
         self.env.init_render(self.ax, self.state_seq[frame], self.obs_seq[frame], lidar=self.plot_lidar, agent=self.plot_agent)