Gym envs: fix observation returning and reset routine

crazyflow/gymnasium_envs/crazyflow.py

@@ -157,23 +157,29 @@ def reset(
         if seed is not None:
             self.jax_key = jax.random.key(seed)
 
-        self.reset_masked(mask=jnp.ones((self.sim.n_worlds), dtype=bool, device=self.device))
+        self.reset_masked(
+            mask=jnp.ones((self.sim.n_worlds), dtype=bool, device=self.device), reset_params=options
+        )
         self.prev_done = jnp.zeros((self.sim.n_worlds), dtype=bool, device=self.device)
         return self._obs(), {}
 
     def reset_masked(self, mask: Array, reset_params: dict | None = None) -> None:
-        default_reset_params = {
-            "pos_min": jnp.array([-1.0, -1.0, 1.0]),  # x,y,z
-            "pos_max": jnp.array([1.0, 1.0, 2.0]),  # x,y,z
-            "vel_min": -1.0,
-            "vel_max": 1.0,
+        if reset_params is None:
+            reset_params = {}
+
+        default_drone_reset_params = {
+            "pos_min": reset_params.pop("pos_min", jnp.array([-1.0, -1.0, 1.0])),  # x,y,z
+            "pos_max": reset_params.pop("pos_max", jnp.array([1.0, 1.0, 2.0])),  # x,y,z
+            "vel_min": reset_params.pop("vel_min", -1.0),
+            "vel_max": reset_params.pop("vel_max", 1.0),
         }
 
-        if reset_params is not None:
-            invalid_keys = set(reset_params.keys()) - set(default_reset_params.keys())
-            if invalid_keys:
-                raise ValueError(f"Invalid bounds keys: {invalid_keys}")
-            default_reset_params.update(reset_params)
+        # sanity check to see if all keys have been used
+        if len(reset_params) > 0:
+            warnings.warn(
+                f"Unused reset parameters: {reset_params.keys()}. "
+                "These will be ignored in the reset function. In case this parameter has already been used, please make sure to pop it from the dictionary."
+            )
 
         self.sim.reset(mask=mask)
         mask3d = mask[:, None, None]
@@ -183,8 +189,8 @@ def reset_masked(self, mask: Array, reset_params: dict | None = None) -> None:
         init_pos = jax.random.uniform(
             key=subkey,
             shape=(self.sim.n_worlds, self.sim.n_drones, 3),
-            minval=default_reset_params["pos_min"],
-            maxval=default_reset_params["pos_max"],
+            minval=default_drone_reset_params["pos_min"],
+            maxval=default_drone_reset_params["pos_max"],
         )
         self.sim.data = self.sim.data.replace(
             states=self.sim.data.states.replace(
@@ -196,8 +202,8 @@ def reset_masked(self, mask: Array, reset_params: dict | None = None) -> None:
         init_vel = jax.random.uniform(
             key=subkey,
             shape=(self.sim.n_worlds, self.sim.n_drones, 3),
-            minval=default_reset_params["vel_min"],
-            maxval=default_reset_params["vel_max"],
+            minval=default_drone_reset_params["vel_min"],
+            maxval=default_drone_reset_params["vel_max"],
         )
         self.sim.data = self.sim.data.replace(
             states=self.sim.data.states.replace(
@@ -242,7 +248,9 @@ def render(self):
     def _obs(self) -> dict[str, Array]:
         fields = self.obs_keys
         states = [getattr(self.sim.data.states, field) for field in fields]
-        return {k: v.squeeze() for k, v in zip(fields, states)}
+        return {
+            k: v[:, 0, :] for k, v in zip(fields, states)
+        }  # drop n_drones dimension, as it is always 1 for now
 
     def close(self):
         self.sim.close()
@@ -273,19 +281,22 @@ def _reward(prev_done: Array, terminated: Array, states: SimState, goal: Array)
         reward = jnp.where(prev_done.reshape(-1, 1), 0.0, reward)
         return reward
 
-    def reset_masked(self, mask: Array) -> None:
-        super().reset_masked(mask)
+    def reset_masked(self, mask: Array, reset_params: dict | None = None) -> None:
+        if reset_params is None:
+            reset_params = {}
 
         # Generate new goals
         self.jax_key, subkey = jax.random.split(self.jax_key)
         new_goals = jax.random.uniform(
             key=subkey,
             shape=(self.sim.n_worlds, 3),
-            minval=jnp.array([-1.0, -1.0, 0.5]),  # x,y,z
-            maxval=jnp.array([1.0, 1.0, 1.5]),  # x,y,z
+            minval=reset_params.pop("goal_pos_min", jnp.array([-1.0, -1.0, 0.5])),  # x,y,z
+            maxval=reset_params.pop("goal_pos_max", jnp.array([1.0, 1.0, 1.5])),  # x,y,z
         )
         self.goal = self.goal.at[mask].set(new_goals[mask])
 
+        super().reset_masked(mask, reset_params)
+
     def step(self, action: Array) -> tuple[Array, Array, Array, Array, dict]:
         if self.render_goal_marker:
             for i in range(self.sim.n_worlds):
@@ -300,7 +311,9 @@ def step(self, action: Array) -> tuple[Array, Array, Array, Array, dict]:
 
     def _obs(self) -> dict[str, Array]:
         obs = super()._obs()
-        obs["difference_to_goal"] = [self.goal - self.sim.data.states.pos]
+        obs["difference_to_goal"] = (
+            self.goal - self.sim.data.states.pos[:, 0, :]
+        )  # drop n_drones dimension, as it is always 1 for now
         return obs
 
 
@@ -329,22 +342,27 @@ def _reward(prev_done: Array, terminated: Array, states: SimState, target_vel: A
         reward = jnp.where(prev_done.reshape(-1, 1), 0.0, reward)
         return reward
 
-    def reset_masked(self, mask: Array) -> None:
-        super().reset_masked(mask)
+    def reset_masked(self, mask: Array, reset_params: dict | None = None) -> None:
+        if reset_params is None:
+            reset_params = {}
 
         # Generate new target_vels
         self.jax_key, subkey = jax.random.split(self.jax_key)
         new_target_vel = jax.random.uniform(
             key=subkey,
             shape=(self.sim.n_worlds, 3),
-            minval=jnp.array([-1.0, -1.0, -1.0]),  # x,y,z
-            maxval=jnp.array([1.0, 1.0, 1.0]),  # x,y,z
+            minval=reset_params.pop("target_vel_min", jnp.array([-1.0, -1.0, -1.0])),  # x,y,z
+            maxval=reset_params.pop("target_vel_max", jnp.array([1.0, 1.0, 1.0])),  # x,y,z
         )
         self.target_vel = self.target_vel.at[mask].set(new_target_vel[mask])
 
+        super().reset_masked(mask)
+
     def _obs(self) -> dict[str, Array]:
         obs = super()._obs()
-        obs["difference_to_target_vel"] = [self.target_vel - self.sim.data.states.vel]
+        obs["difference_to_target_vel"] = (
+            self.target_vel - self.sim.data.states.vel[:, 0, :]
+        )  # drop n_drones dimension, as it is always 1 for now
         return obs
 
 
@@ -375,9 +393,6 @@ def _reward(prev_done: Array, terminated: Array, states: SimState, goal: Array)
         reward = jnp.where(prev_done.reshape(-1, 1), 0.0, reward)
         return reward
 
-    def reset_masked(self, mask: Array) -> None:
-        super().reset_masked(mask)
-
     def step(self, action: Array) -> tuple[Array, Array, Array, Array, dict]:
         if self.render_landing_target:
             for i in range(self.sim.n_worlds):
@@ -392,7 +407,9 @@ def step(self, action: Array) -> tuple[Array, Array, Array, Array, dict]:
 
     def _obs(self) -> dict[str, Array]:
         obs = super()._obs()
-        obs["difference_to_goal"] = [self.goal - self.sim.data.states.pos]
+        obs["difference_to_goal"] = (
+            self.goal - self.sim.data.states.pos[:, 0, :]
+        )  # drop n_drones dimension, as it is always 1 for now
         return obs
 
 
@@ -478,14 +495,19 @@ def _reward(prev_done: Array, terminated: Array, states: SimState, goal: Array)
         reward = jnp.where(prev_done.reshape(-1, 1), 0.0, reward)
         return reward
 
-    def reset_masked(self, mask: Array) -> None:
-        reset_params = {
-            "pos_min": jnp.array([-0.1, -0.1, 1.1]),  # x,y,z
-            "pos_max": jnp.array([0.1, 0.1, 1.3]),  # x,y,z
-            "vel_min": -0.5,
-            "vel_max": 0.5,
+    def reset_masked(self, mask: Array, reset_params: dict | None = None) -> None:
+        if reset_params is None:
+            reset_params = {}
+
+        # Different initial conditions than CrazyflowBaseEnv
+        default_drone_reset_params = {
+            "pos_min": reset_params.pop("pos_min", jnp.array([-0.1, -0.1, 1.1])),  # x,y,z
+            "pos_max": reset_params.pop("pos_max", jnp.array([0.1, 0.1, 1.3])),  # x,y,z
+            "vel_min": reset_params.pop("vel_min", -0.5),
+            "vel_max": reset_params.pop("vel_max", 0.5),
         }
-        super().reset_masked(mask, reset_params)
+
+        super().reset_masked(mask, default_drone_reset_params)
 
     def _obs(self) -> dict[str, Array]:
         obs = super()._obs()

examples/gymnasium_env.py

@@ -9,17 +9,27 @@
 def main():
     enable_cache()
     SEED = 42
-    envs = gymnasium.make_vec("DroneLanding-v0", num_envs=20, freq=50, time_horizon_in_seconds=2)
+    envs = gymnasium.make_vec("DroneReachPos-v0", num_envs=20, freq=50, time_horizon_in_seconds=2)
 
-    # This wrapper makes it possible to interact with the environment using numpy arrays, if
-    # desired. JaxToTorch is available as well.
+    # This wrapper makes it possible to interact with the environment using numpy arrays, if desired. JaxToTorch is available as well.
     envs = JaxToNumpy(envs)
 
-    # dummy action for going up (in attitude control)
+    # Dummy action for going up (in attitude control)
     action = np.zeros((20, 4), dtype=np.float32)
     action[..., 0] = 0.4
 
-    obs, info = envs.reset(seed=SEED)
+    # Environments provide reset parameters that can be used to set the initial state of the environment.
+    obs, info = envs.reset(
+        seed=SEED,
+        options={
+            "pos_min": np.array([-1.0, 1.0, 1.0]),
+            "pos_max": np.array([-1.0, 1.0, 1.0]),
+            "vel_min": 0.0,
+            "vel_max": 0.0,
+            "goal_pos_min": np.array([-1.0, 1.0, 1.0]),
+            "goal_pos_max": np.array([-1.0, 1.0, 1.0]),
+        },
+    )
 
     # Step through the environment
     for _ in range(100):

tests/integration/test_gymnasium_envs.py

@@ -0,0 +1,29 @@
+import gymnasium
+import numpy as np
+import pytest
+from gymnasium.wrappers.vector import JaxToNumpy
+
+import crazyflow  # noqa: F401, register gymnasium envs
+
+
+@pytest.mark.integration
+def test_gymnasium_reset():
+    """Test reset behavior of the DroneReachPos-v0 environment."""
+    SEED = 42
+    envs = gymnasium.make_vec("DroneReachPos-v0", num_envs=1, freq=50, time_horizon_in_seconds=2)
+
+    envs = JaxToNumpy(envs)
+    obs, _ = envs.reset(
+        seed=SEED,
+        options={
+            "pos_min": np.array([-1.0, 1.0, 1.0]),
+            "pos_max": np.array([-1.0, 1.0, 1.0]),
+            "vel_min": 0.0,
+            "vel_max": 0.0,
+            "goal_pos_min": np.array([-1.0, 1.0, 1.0]),
+            "goal_pos_max": np.array([-1.0, 1.0, 1.0]),
+        },
+    )
+    assert np.all(obs["pos"] == np.array([[-1.0, 1.0, 1.0]]))
+    assert np.all(obs["difference_to_goal"] == np.array([[.0, .0, .0]]))
+    assert np.all(obs["vel"] == np.array([[0.0, 0.0, 0.0]]))

tests/integration/test_reset.py

@@ -2,6 +2,7 @@
 import numpy as np
 import pytest
 
+import crazyflow  # noqa: F401, register gymnasium envs
 from crazyflow.control import Control
 from crazyflow.sim import Physics, Sim
 
@@ -62,3 +63,4 @@ def test_reset_multi_world(physics: Physics):
         sim.step(sim.freq // sim.control_freq)
     assert jnp.all(sim.data.states.pos == final_pos)
     assert jnp.all(sim.data.states.quat == final_quat)
+