Fix gymnasium examples

Author: amacati

examples/gymnasium_env.py

@@ -1,40 +1,28 @@
 import gymnasium
 import numpy as np
 from gymnasium.wrappers.vector import JaxToNumpy  # , JaxToTorch
-from ml_collections import config_dict
 
-from crazyflow.control import Control
-from crazyflow.sim.physics import Physics
+import crazyflow  # noqa: F401, register gymnasium envs
+from crazyflow.utils import enable_cache
 
 
 def main():
-    # set config for simulation
-    sim_config = config_dict.ConfigDict()
-    sim_config.device = "cpu"
-    sim_config.physics = Physics.sys_id
-    sim_config.control = Control.attitude
-    sim_config.attitude_freq = 50
-    sim_config.n_drones = 1
-    sim_config.n_worlds = 20
-
+    enable_cache()
     SEED = 42
-
-    envs = gymnasium.make_vec(
-        "DroneLanding-v0", time_horizon_in_seconds=2, num_envs=sim_config.n_worlds, **sim_config
-    )
+    envs = gymnasium.make_vec("DroneLanding-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.
     envs = JaxToNumpy(envs)
 
     # dummy action for going up (in attitude control)
-    action = np.zeros((sim_config.n_worlds * sim_config.n_drones, 4), dtype=np.float32)
+    action = np.zeros((20, 4), dtype=np.float32)
     action[..., 0] = 0.4
 
     obs, info = envs.reset(seed=SEED)
 
     # Step through the environment
-    for _ in range(1500):
+    for _ in range(100):
         observation, reward, terminated, truncated, info = envs.step(action)
         envs.render()
 

examples/gymnasium_env_trajectory.py

@@ -1,38 +1,25 @@
 import gymnasium
 import numpy as np
 from gymnasium.wrappers.vector import JaxToNumpy  # , JaxToTorch
-from ml_collections import config_dict
-from scipy.interpolate import splev
 
-from crazyflow.control import Control
 from crazyflow.gymnasium_envs import CrazyflowRL  # noqa: F401
-from crazyflow.sim.physics import Physics
+from crazyflow.utils import enable_cache
 
 
 def main():
-    # set config for simulation
-    sim_config = config_dict.ConfigDict()
-    sim_config.device = "cpu"
-    sim_config.physics = Physics.sys_id
-    sim_config.control = Control.attitude
-    sim_config.attitude_freq = 50
-    sim_config.n_drones = 1
-    sim_config.n_worlds = 20
-
+    enable_cache()
     SEED = 42
-
     # Create environment that contains a figure eight trajectory. You can parametrize the
     # observation space, i.e., which part of the trajectory is contained in the observation. Please
     # refer to the documentation of the environment for more information.
     envs = gymnasium.make_vec(
         "DroneFigureEightTrajectory-v0",
-        n_trajectory_sample_points=10,
-        dt_trajectory_sample_points=0.1,
+        num_envs=20,
+        freq=50,
+        n_samples=10,
+        samples_dt=0.1,
         trajectory_time=10.0,
-        render_trajectory_sample=True,  # useful for debug purposes
-        time_horizon_in_seconds=10.0,
-        num_envs=sim_config.n_worlds,
-        **sim_config,
+        render_samples=True,
     )
 
     # RL wrapper to clip the actions to [-1, 1] and rescale them for use with common DRL libraries.
@@ -43,19 +30,12 @@ def main():
     envs = JaxToNumpy(envs)
 
     # dummy action for going up (in attitude control)
-    action = np.zeros((sim_config.n_worlds * sim_config.n_drones, 4), dtype=np.float32)
-    action[..., 0] = 0.34
+    action = np.zeros((20, 4), dtype=np.float32)
+    action[..., 0] = 0.31
 
     obs, info = envs.reset(seed=SEED)
-
-    # The trajectory is defined as a scipy spline. Its parameter can be retrieved using
-    # `envs.unwrapped.tck`. The spline can be reconstructed using scipy's splev.
-    spline_params = envs.unwrapped.tck
-    tau = envs.unwrapped.tau  # 1D parameters of the spline for the current timestep, in [0,1]
-    value = splev(tau, spline_params)  # noqa: F841, used for demonstration purposes
-
     # Step through the environment
-    for _ in range(1500):
+    for _ in range(500):
         observation, reward, terminated, truncated, info = envs.step(action)
         envs.render()