Improve build interface

Author: amacati

crazyflow/sim/sim.py

@@ -75,7 +75,8 @@ def __init__(
         self.viewer: MujocoRenderer | None = None
 
         self.data = self.init_data(state_freq, attitude_freq, thrust_freq, rng_key, mjx_data)
-        self.default_data = self.init_default_data()
+        self.default_data: SimData
+        self.build_default_data()
 
         # Build the simulation pipeline and overwrite the default _step implementation with it
         self.reset_pipeline: tuple[Callable[[SimData, Array[bool] | None], SimData], ...] = tuple()
@@ -96,6 +97,88 @@ def __init__(
         self.build_reset_fn()
         self.build_step_fn()
 
+    def reset(self, mask: Array | None = None):
+        """Reset the simulation to the initial state.
+
+        Args:
+            mask: Boolean array of shape (n_worlds, ) that indicates which worlds to reset. If None,
+                all worlds are reset.
+        """
+        assert mask is None or mask.shape == (self.n_worlds,), f"Mask shape mismatch {mask.shape}"
+        self.data = self._reset(self.data, self.default_data, mask)
+
+    def step(self, n_steps: int = 1):
+        """Simulate all drones in all worlds for n time steps."""
+        assert n_steps > 0, "Number of steps must be positive"
+        self.data = self._step(self.data, n_steps=n_steps)
+
+    def attitude_control(self, controls: Array):
+        """Set the desired attitude for all drones in all worlds.
+
+        We need to stage the attitude controls because the sys_id physics mode operates directly on
+        the attitude controls. If we were to directly update the controls, this would effectively
+        bypass the control frequency and run the attitude controller at the physics update rate. By
+        staging the controls, we ensure that the physics module sees the old controls until the
+        controller updates at its correct frequency.
+        """
+        assert controls.shape == (self.n_worlds, self.n_drones, 4), "controls shape mismatch"
+        assert self.control == Control.attitude, "Attitude control is not enabled by the sim config"
+        controls = to_device(controls, self.device)
+        self.data = self.data.replace(controls=self.data.controls.replace(staged_attitude=controls))
+
+    def state_control(self, controls: Array):
+        """Set the desired state for all drones in all worlds."""
+        assert controls.shape == (self.n_worlds, self.n_drones, 13), "controls shape mismatch"
+        assert self.control == Control.state, "State control is not enabled by the sim config"
+        controls = to_device(controls, self.device)
+        self.data = self.data.replace(controls=self.data.controls.replace(state=controls))
+
+    def thrust_control(self, cmd: Array):
+        """Set the desired thrust for all drones in all worlds."""
+        assert cmd.shape == (self.n_worlds, self.n_drones, 4), "Command shape mismatch"
+        assert self.control == Control.thrust, "Thrust control is not enabled by the sim config"
+        controls = to_device(cmd, self.device)
+        self.data = self.data.replace(controls=self.data.controls.replace(thrust=controls))
+
+    def render(
+        self,
+        mode: str | None = "human",
+        world: int = 0,
+        default_cam_config: dict | None = None,
+        width: int = 640,
+        height: int = 480,
+    ) -> NDArray | None:
+        if self.viewer is None:
+            patch_viewer()
+            self.mj_model.vis.global_.offwidth = width
+            self.mj_model.vis.global_.offheight = height
+            self.viewer = MujocoRenderer(
+                self.mj_model,
+                self.mj_data,
+                max_geom=self.max_visual_geom,
+                default_cam_config=default_cam_config,
+                height=height,
+                width=width,
+            )
+        self.mj_data.qpos[:] = self.data.mjx_data.qpos[world, :]
+        self.mj_data.mocap_pos[:] = self.data.mjx_data.mocap_pos[world, :]
+        self.mj_data.mocap_quat[:] = self.data.mjx_data.mocap_quat[world, :]
+        mujoco.mj_forward(self.mj_model, self.mj_data)
+        return self.viewer.render(mode)
+
+    def seed(self, seed: int):
+        """Set the JAX rng key for the simulation.
+
+        Args:
+            seed: The seed for the JAX rng.
+        """
+        self.data = seed_sim(self.data, seed, self.device)
+
+    def close(self):
+        if self.viewer is not None:
+            self.viewer.close()
+        self.viewer = None
+
     def build_mjx_spec(self) -> mujoco.MjSpec:
         """Build the MuJoCo model specification for the simulation."""
         assert self._xml_path.exists(), f"Model file {self._xml_path} does not exist"
@@ -167,6 +250,42 @@ def step(data: SimData, n_steps: int = 1) -> SimData:
 
         self._step = step
 
+    def build_reset_fn(self):
+        """Build the reset function for the current simulation configuration."""
+        pipeline = self.reset_pipeline
+
+        @jax.jit
+        def reset(data: SimData, default_data: SimData, mask: Array | None = None) -> SimData:
+            data = pytree_replace(data, default_data, mask)  # Does not overwrite rng_key
+            for fn in pipeline:
+                data = fn(data, mask)
+            data = self.sync_sim2mjx(data, self.mjx_model)
+            return data
+
+        self._reset = reset
+
+    def build_data(self):
+        self.data = self.init_data(
+            self.data.controls.state_freq,
+            self.data.controls.attitude_freq,
+            self.data.controls.thrust_freq,
+            self.data.core.rng_key,
+            self.data.mjx_data,
+        )
+
+    def build_default_data(self):
+        """Initialize the default data for the simulation."""
+        self.default_data = self.data.replace()
+
+    def build_mjx(self):
+        if self.viewer is not None:
+            self.viewer.close()
+            self.viewer = None
+        self.mj_model, self.mj_data, self.mjx_model, mjx_data = self.build_mjx_model(self.spec)
+        self.data = self.data.replace(mjx_data=mjx_data)
+        self.data = self.sync_sim2mjx(self.data, self.mjx_model)
+        self.default_data = self.default_data.replace(mjx_data=mjx_data)
+
     def init_data(
         self, state_freq: int, attitude_freq: int, thrust_freq: int, rng_key: Array, mjx_data: Data
     ) -> tuple[SimData, SimData]:
@@ -189,161 +308,6 @@ def init_data(
         data = self.sync_sim2mjx(data, self.mjx_model)
         return data
 
-    def init_default_data(self) -> SimData:
-        """Initialize the default data for the simulation.
-
-        Todo:
-            Only save the data of one world.
-        """
-        return self.data.replace()
-
-    def build_reset_fn(self):
-        """Build the reset function for the current simulation configuration."""
-        pipeline = self.reset_pipeline
-
-        @jax.jit
-        def reset(data: SimData, default_data: SimData, mask: Array | None = None) -> SimData:
-            data = pytree_replace(data, default_data, mask)  # Does not overwrite rng_key
-            for fn in pipeline:
-                data = fn(data, mask)
-            data = self.sync_sim2mjx(data, self.mjx_model)
-            return data
-
-        self._reset = reset
-
-    def build(
-        self,
-        *,
-        mjx: bool = True,
-        data: bool = True,
-        default_data: bool = True,
-        reset: bool = True,
-        step: bool = True,
-    ):
-        """Build the simulation pipeline.
-
-        This method is used to (re)build the simulation pipeline after changing the MuJoCo
-        model specification or any of the default functions that are used in the compiled step
-        function.
-
-        Warning:
-            Depending on what you build, you reset the simulation state. For example, rebuilding the
-            simulation data will reset the drone states.
-
-        Args:
-            mjx: Flag to (re)build the MuJoCo model and data structures.
-            data: Flag to (re)build the simulation data.
-            default_data: Flag to (re)build the default data. Useful for setting the reset state to
-                the current state.
-            reset: Flag to (re)build the reset function.
-            step: Flag to (re)build the simulation step function.
-        """
-        # TODO: Write tests for all options
-        if mjx:
-            if self.viewer is not None:
-                self.viewer.close()
-                self.viewer = None
-            self.mj_model, self.mj_data, self.mjx_model, mjx_data = self.build_mjx_model(self.spec)
-            self.data = self.data.replace(mjx_data=mjx_data)
-            self.data = self.sync_sim2mjx(self.data, self.mjx_model)
-            self.default_data = self.default_data.replace(mjx_data=mjx_data)
-        if data:
-            self.data = self.init_data(
-                self.data.controls.state_freq,
-                self.data.controls.attitude_freq,
-                self.data.controls.thrust_freq,
-                self.data.core.rng_key,
-                self.data.mjx_data if not mjx else mjx_data,
-            )
-        if default_data:
-            self.default_data = self.init_default_data()
-        if reset:
-            self.build_reset_fn()
-        if step:
-            self.build_step_fn()
-
-    def reset(self, mask: Array | None = None):
-        """Reset the simulation to the initial state.
-
-        Args:
-            mask: Boolean array of shape (n_worlds, ) that indicates which worlds to reset. If None,
-                all worlds are reset.
-        """
-        assert mask is None or mask.shape == (self.n_worlds,), f"Mask shape mismatch {mask.shape}"
-        self.data = self._reset(self.data, self.default_data, mask)
-
-    def step(self, n_steps: int = 1):
-        """Simulate all drones in all worlds for n time steps."""
-        assert n_steps > 0, "Number of steps must be positive"
-        self.data = self._step(self.data, n_steps=n_steps)
-
-    def attitude_control(self, controls: Array):
-        """Set the desired attitude for all drones in all worlds.
-
-        We need to stage the attitude controls because the sys_id physics mode operates directly on
-        the attitude controls. If we were to directly update the controls, this would effectively
-        bypass the control frequency and run the attitude controller at the physics update rate. By
-        staging the controls, we ensure that the physics module sees the old controls until the
-        controller updates at its correct frequency.
-        """
-        assert controls.shape == (self.n_worlds, self.n_drones, 4), "controls shape mismatch"
-        assert self.control == Control.attitude, "Attitude control is not enabled by the sim config"
-        controls = to_device(controls, self.device)
-        self.data = self.data.replace(controls=self.data.controls.replace(staged_attitude=controls))
-
-    def state_control(self, controls: Array):
-        """Set the desired state for all drones in all worlds."""
-        assert controls.shape == (self.n_worlds, self.n_drones, 13), "controls shape mismatch"
-        assert self.control == Control.state, "State control is not enabled by the sim config"
-        controls = to_device(controls, self.device)
-        self.data = self.data.replace(controls=self.data.controls.replace(state=controls))
-
-    def thrust_control(self, cmd: Array):
-        """Set the desired thrust for all drones in all worlds."""
-        assert cmd.shape == (self.n_worlds, self.n_drones, 4), "Command shape mismatch"
-        assert self.control == Control.thrust, "Thrust control is not enabled by the sim config"
-        controls = to_device(cmd, self.device)
-        self.data = self.data.replace(controls=self.data.controls.replace(thrust=controls))
-
-    def render(
-        self,
-        mode: str | None = "human",
-        world: int = 0,
-        default_cam_config: dict | None = None,
-        width: int = 640,
-        height: int = 480,
-    ) -> NDArray | None:
-        if self.viewer is None:
-            patch_viewer()
-            self.mj_model.vis.global_.offwidth = width
-            self.mj_model.vis.global_.offheight = height
-            self.viewer = MujocoRenderer(
-                self.mj_model,
-                self.mj_data,
-                max_geom=self.max_visual_geom,
-                default_cam_config=default_cam_config,
-                height=height,
-                width=width,
-            )
-        self.mj_data.qpos[:] = self.data.mjx_data.qpos[world, :]
-        self.mj_data.mocap_pos[:] = self.data.mjx_data.mocap_pos[world, :]
-        self.mj_data.mocap_quat[:] = self.data.mjx_data.mocap_quat[world, :]
-        mujoco.mj_forward(self.mj_model, self.mj_data)
-        return self.viewer.render(mode)
-
-    def seed(self, seed: int):
-        """Set the JAX rng key for the simulation.
-
-        Args:
-            seed: The seed for the JAX rng.
-        """
-        self.data = seed_sim(self.data, seed, self.device)
-
-    def close(self):
-        if self.viewer is not None:
-            self.viewer.close()
-        self.viewer = None
-
     @property
     def time(self) -> Array:
         return self.data.core.steps / self.data.core.freq

examples/disturbance.py

@@ -39,7 +39,7 @@ def main(plot: bool = False):
     # inspect the step pipeline with
     # print(sim.step_pipeline)
     sim.step_pipeline = sim.step_pipeline[:2] + (disturbance_fn,) + sim.step_pipeline[2:]
-    sim.build(mjx=False, data=False, default_data=False, step=True)
+    sim.build_step_fn()
     pos_disturbed, rpy_disturbed = [], []
     sim.reset()
     for _ in range(3 * sim.control_freq):

tests/integration/test_disturbance.py

@@ -29,7 +29,7 @@ def test_disturbance(physics: Physics):
 
     sim.reset()
     sim.step_pipeline = sim.step_pipeline[:2] + (disturbance_fn,) + sim.step_pipeline[2:]
-    sim.build(mjx=False, data=False, default_data=False)
+    sim.build_step_fn()
     for _ in range(sim.control_freq):
         sim.state_control(control)
         sim.step(sim.freq // sim.control_freq)