Merge branch 'main' into dev

Author: amacati

.gitignore

@@ -7,11 +7,11 @@ saves
 .pytest_cache
 *.json
 **/*.csv
+.vscode
 build
 .venv
 !/.devcontainer/devcontainer.json
 !/.devcontainer/devcontainer.linux.json
 !/.devcontainer/devcontainer.wsl2.json
-!/.vscode/launch.json
 **/*.pt
 tutorials/ppo/wandb

.vscode/launch.json

@@ -43,7 +43,12 @@ def profile_gym_env_step(sim_config: config_dict.ConfigDict, n_steps: int, devic
     device = jax.devices(device)[0]
 
     envs = gymnasium.make_vec(
-        "DroneReachPos-v0", time_horizon_in_seconds=3, num_envs=sim_config.n_worlds, **sim_config
+        "DroneReachPos-v0",
+        time_horizon_in_seconds=3,
+        num_envs=sim_config.n_worlds,
+        device=sim_config.device,
+        freq=sim_config.attitude_freq,
+        physics=sim_config.physics,
     )
 
     # Action for going up (in attitude control)

benchmark/main.py

@@ -82,7 +82,7 @@ def _integrate(
     dvel: Array,
     dang_vel: Array,
     dt: float,
-) -> SimData:
+) -> tuple[Array, Array, Array, Array]:
     """Integrate the dynamics forward in time.
 
     Args: