From 4587b5ca89ebfd5922362fff3adde6c6760db3da Mon Sep 17 00:00:00 2001 From: OliEfr Date: Mon, 26 May 2025 18:14:43 +0200 Subject: [PATCH 1/3] Gym envs: fix observation returning and reset routine --- crazyflow/gymnasium_envs/crazyflow.py | 98 ++++++++++++++++----------- examples/gymnasium_env.py | 20 ++++-- tests/integration/test_reset.py | 29 ++++++++ 3 files changed, 104 insertions(+), 43 deletions(-) diff --git a/crazyflow/gymnasium_envs/crazyflow.py b/crazyflow/gymnasium_envs/crazyflow.py index 0c230f4..7148ae4 100644 --- a/crazyflow/gymnasium_envs/crazyflow.py +++ b/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(), {} + return self._obs(), None 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() diff --git a/examples/gymnasium_env.py b/examples/gymnasium_env.py index de9f9f4..3139a6b 100644 --- a/examples/gymnasium_env.py +++ b/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): diff --git a/tests/integration/test_reset.py b/tests/integration/test_reset.py index 941b5a7..76b4a85 100644 --- a/tests/integration/test_reset.py +++ b/tests/integration/test_reset.py @@ -6,6 +6,12 @@ from crazyflow.sim import Physics, Sim +import gymnasium +import numpy as np +from gymnasium.wrappers.vector import JaxToNumpy # , JaxToTorch +import crazyflow # noqa: F401, register gymnasium envs + + @pytest.mark.integration @pytest.mark.parametrize("physics", Physics) def test_reset_during_simulation(physics: Physics): @@ -62,3 +68,26 @@ 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) + +@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]])) + \ No newline at end of file From d46892b4c0b0e0b25c631e6182acc2816e6ca973 Mon Sep 17 00:00:00 2001 From: OliEfr Date: Mon, 26 May 2025 20:04:46 +0200 Subject: [PATCH 2/3] Implement feedback --- crazyflow/gymnasium_envs/crazyflow.py | 2 +- tests/integration/test_gymnasium_envs.py | 29 ++++++++++++++++++++++ tests/integration/test_reset.py | 31 +++--------------------- 3 files changed, 33 insertions(+), 29 deletions(-) create mode 100644 tests/integration/test_gymnasium_envs.py diff --git a/crazyflow/gymnasium_envs/crazyflow.py b/crazyflow/gymnasium_envs/crazyflow.py index 7148ae4..2139164 100644 --- a/crazyflow/gymnasium_envs/crazyflow.py +++ b/crazyflow/gymnasium_envs/crazyflow.py @@ -161,7 +161,7 @@ def reset( 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(), None + return self._obs(), {} def reset_masked(self, mask: Array, reset_params: dict | None = None) -> None: if reset_params is None: diff --git a/tests/integration/test_gymnasium_envs.py b/tests/integration/test_gymnasium_envs.py new file mode 100644 index 0000000..c451c58 --- /dev/null +++ b/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]])) \ No newline at end of file diff --git a/tests/integration/test_reset.py b/tests/integration/test_reset.py index 76b4a85..49c9c7d 100644 --- a/tests/integration/test_reset.py +++ b/tests/integration/test_reset.py @@ -1,17 +1,14 @@ +import gymnasium import jax.numpy as jnp import numpy as np import pytest +from gymnasium.wrappers.vector import JaxToNumpy # , JaxToTorch +import crazyflow # noqa: F401, register gymnasium envs from crazyflow.control import Control from crazyflow.sim import Physics, Sim -import gymnasium -import numpy as np -from gymnasium.wrappers.vector import JaxToNumpy # , JaxToTorch -import crazyflow # noqa: F401, register gymnasium envs - - @pytest.mark.integration @pytest.mark.parametrize("physics", Physics) def test_reset_during_simulation(physics: Physics): @@ -68,26 +65,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) - -@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]])) \ No newline at end of file From 4225dba27583aa0d71fadcd63d06df26cc4b7d0d Mon Sep 17 00:00:00 2001 From: OliEfr Date: Mon, 26 May 2025 20:05:40 +0200 Subject: [PATCH 3/3] Implement feedback --- tests/integration/test_reset.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/tests/integration/test_reset.py b/tests/integration/test_reset.py index 49c9c7d..8e7a368 100644 --- a/tests/integration/test_reset.py +++ b/tests/integration/test_reset.py @@ -1,8 +1,6 @@ -import gymnasium import jax.numpy as jnp import numpy as np import pytest -from gymnasium.wrappers.vector import JaxToNumpy # , JaxToTorch import crazyflow # noqa: F401, register gymnasium envs from crazyflow.control import Control