Add functional API. Include formatting check in CI

Author: amacati

.github/workflows/ruff.yml

@@ -1,8 +1,16 @@
 name: Ruff
 on: [ push, pull_request ]
 jobs:
-  ruff:
+  ruff-check:
     runs-on: ubuntu-latest
     steps:
       - uses: actions/checkout@v4
-      - uses: astral-sh/ruff-action@v1
+      - uses: astral-sh/ruff-action@v1
+
+  ruff-format:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      - uses: astral-sh/ruff-action@v1
+        with:
+          args: "format --check --diff"

crazyflow/sim/sim.py

@@ -98,8 +98,7 @@ def __init__(
         self.viewer: MujocoRenderer | None = None
 
         self.data = self.init_data(state_freq, attitude_freq, force_torque_freq, rng_key)
-        self.default_data: SimData
-        self.build_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()
@@ -115,8 +114,8 @@ def __init__(
         # enable checks for negative z sign
         self.step_pipeline += (clip_floor_pos,)
 
-        self.build_reset_fn()
-        self.build_step_fn()
+        self._reset = self.build_reset_fn()
+        self._step = self.build_step_fn()
 
     def reset(self, mask: Array | None = None):
         """Reset the simulation to the initial state.
@@ -251,15 +250,23 @@ def build_mjx_model(self, spec: mujoco.MjSpec) -> tuple[Any, Any, Model, Data]:
         mjx_data = jax.vmap(lambda _: mjx_data)(jnp.arange(self.n_worlds))
         return mj_model, mj_data, mjx_model, mjx_data
 
-    def build_step_fn(self):
+    def build_step_fn(self) -> Callable[[SimData, int], SimData]:
         """Setup the chain of functions that are called in Sim.step().
 
         We know all the functions that are called in succession since the simulation is configured
         at initialization time. Instead of branching through options at runtime, we construct a step
         function at initialization that selects the correct functions based on the settings.
 
+        Note:
+            This function both changes the underlying implementation of Sim.step() in-place to the
+            current pipeline and returns the function for pure functional style programming.
+
         Warning:
             If any settings change, the pipeline of functions needs to be reconstructed.
+
+        Returns:
+            The pure JAX function that steps through the simulation. It takes the current SimData
+            and the number of steps to simulate, and returns the updated SimData.
         """
         pipeline = self.step_pipeline
 
@@ -271,9 +278,9 @@ def single_step(data: SimData, _: None) -> tuple[SimData, None]:
 
         # ``scan`` allows us control over loop unrolling for single steps from a single WhileOp to
         # complete unrolling, reducing either compilation times or fusing the loops to give XLA
-        # maximum freedom to reorder operations and jointly optimize the pipeline. This is especially
-        # relevant for the common use case of running multiple sim steps in an outer loop, e.g. in
-        # gym environments.
+        # maximum freedom to reorder operations and jointly optimize the pipeline. This is
+        # especially relevant for the common use case of running multiple sim steps in an outer
+        # loop, e.g. in gym environments.
         # Having n_steps as a static argument is fine, since patterns with n_steps > 1 will almost
         # always use the same n_steps value for successive calls.
         @partial(jax.jit, static_argnames="n_steps")
@@ -283,9 +290,19 @@ def step(data: SimData, n_steps: int = 1) -> SimData:
             return data
 
         self._step = step
+        return step
+
+    def build_reset_fn(self) -> Callable[[SimData, SimData, Array | None], SimData]:
+        """Build the reset function for the current simulation configuration.
 
-    def build_reset_fn(self):
-        """Build the reset function for the current simulation configuration."""
+        Note:
+            This function both changes the underlying implementation of Sim.reset() in-place to the
+            current pipeline and returns the function for pure functional style programming.
+
+        Returns:
+            The pure JAX function that resets simulation data. It takes the current SimData, default
+            SimData, and an optional mask for worlds to reset, returning the updated SimData.
+        """
         pipeline = self.reset_pipeline
 
         @jax.jit
@@ -297,18 +314,43 @@ def reset(data: SimData, default_data: SimData, mask: Array | None = None) -> Si
             return data
 
         self._reset = reset
+        return reset
+
+    def build_data(self) -> SimData:
+        """Build the simulation data for the current configuration.
 
-    def build_data(self):
+        Note:
+            This function re-initializes the simulation data according to the current configuration.
+            It also returns the constructed data for use with pure functions.
+
+        Returns:
+            The simulation data as a single PyTree that can be passed to the pure simulation
+            functions for stepping and resetting.
+        """
+        state_freq = self.data.controls.state.freq if self.data.controls.state is not None else 0
+        attitude_freq = (
+            self.data.controls.attitude.freq if self.data.controls.attitude is not None else 0
+        )
+        force_torque_freq = self.data.controls.force_torque.freq
         self.data = self.init_data(
-            self.data.controls.state_freq,
-            self.data.controls.attitude_freq,
-            self.data.controls.force_torque_freq,
-            self.data.core.rng_key,
+            state_freq, attitude_freq, force_torque_freq, self.data.core.rng_key
         )
+        return self.data
 
-    def build_default_data(self):
-        """Initialize the default data for the simulation."""
+    def build_default_data(self) -> SimData:
+        """Initialize the default data for the simulation.
+
+        Note:
+            This function initializes the default data used as a reference in the reset function to
+            reset the simulation to. It also returns the constructed data for use with pure
+            functions.
+
+        Returns:
+            The default simulation data used as a reference in the reset function to reset the
+            simulation to.
+        """
         self.default_data = self.data.replace()
+        return self.default_data
 
     def build_mjx(self):
         if self.viewer is not None:

examples/gymnasium_env.py

@@ -19,13 +19,10 @@ def main():
     action = np.zeros((20, 4), dtype=np.float32)
     action[..., 0] = 0.4
 
-    # Environments provide reset parameters that can be used to set the initial state of the environment.
+    # Environments provide reset parameters that can be used to set the initial state of the
+    # environment.
     obs, info = envs.reset(
         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]),
         }

pixi.lock

@@ -81,11 +81,14 @@ indent-width = 4
 target-version = "py312"
 
 [tool.ruff.lint]
-select = ["E4", "E7", "E9", "F", "I", "D", "TCH", "ANN"]
+select = ["E", "F", "I", "D", "TCH", "ANN"]
 ignore = ["ANN401"]
 fixable = ["ALL"]
 unfixable = []
 
+[tool.ruff.lint.isort]  # Prevent ruff from reporting conflicting settings with isort
+split-on-trailing-comma = false
+
 [tool.ruff.lint.per-file-ignores]
 "benchmark/*" = ["D100", "D103"]
 "tests/*" = ["D100", "D103", "D104"]

pyproject.toml

@@ -15,7 +15,7 @@
 from crazyflow.control import Control
 from crazyflow.exception import ConfigError
 from crazyflow.sim import Physics, Sim
-from crazyflow.sim.data import ControlData
+from crazyflow.sim.data import ControlData, SimData
 from crazyflow.sim.sim import sync_sim2mjx
 from crazyflow.sim.visualize import change_material
 
@@ -534,24 +534,43 @@ def test_change_material_errors(device: str):
     emission = np.ones((n_drones,), dtype=float)
 
     with pytest.raises(ValueError):
-        change_material(
-            sim, mat_name="bad_mat", drone_ids=drone_ids, rgba=rgba, emission=emission
-        )
+        change_material(sim, mat_name="bad_mat", drone_ids=drone_ids, rgba=rgba, emission=emission)
 
     with pytest.raises(ValueError, match="drone_ids must be 1D array"):
         change_material(
-            sim,
-            mat_name="led_top",
-            drone_ids=np.array(2, dtype=int),
-            rgba=rgba,
-            emission=emission,
+            sim, mat_name="led_top", drone_ids=np.array(2, dtype=int), rgba=rgba, emission=emission
         )
 
     with pytest.raises(ValueError, match=r"drone_ids must be in range \[0, 1\]"):
         change_material(
-            sim,
-            mat_name="led_top",
-            drone_ids=np.arange(3, dtype=int),
-            rgba=rgba,
-            emission=emission,
+            sim, mat_name="led_top", drone_ids=np.arange(3, dtype=int), rgba=rgba, emission=emission
         )
+
+
+@pytest.mark.unit
+@pytest.mark.parametrize("control", Control)
+def test_build_data(control: Control):
+    sim = Sim(control=control)
+    data = sim.build_data()
+    assert isinstance(data, SimData), "build_data() must return a SimData instance"
+    default_data = sim.build_default_data()
+    assert isinstance(default_data, SimData), "build_default_data() must return a SimData instance"
+
+
+@pytest.mark.unit
+def test_functional_api():
+    """Test that the functional API works as expected."""
+    sim = Sim()
+    reset_fn = sim.build_reset_fn()
+    step_fn = sim.build_step_fn()
+    # Test types
+    assert callable(reset_fn), "reset_fn must be a pure function"
+    assert not hasattr(reset_fn, "__self__"), "reset_fn must not be a bound method"
+    assert callable(step_fn), "step_fn must be a pure function"
+    assert not hasattr(step_fn, "__self__"), "step_fn must not be a bound method"
+    # Test the functions run as expected
+    data, default_data = sim.build_data(), sim.build_default_data()
+    data = reset_fn(data, default_data, None)
+    data = reset_fn(data, default_data, jnp.array([True] * sim.n_worlds))
+    data = step_fn(data, 1)
+    data = step_fn(data, 2)

tests/unit/test_sim.py

@@ -27,16 +27,16 @@ def test_enable_cache(enable_xla: bool):
         )
 
         assert cache_path == jax.config.jax_compilation_cache_dir, "Cache path not set correctly"
-        assert (
-            min_size == jax.config.jax_persistent_cache_min_entry_size_bytes
-        ), "Min size not set correctly"
-        assert (
-            min_time == jax.config.jax_persistent_cache_min_compile_time_secs
-        ), "Min time not set correctly"
+        assert min_size == jax.config.jax_persistent_cache_min_entry_size_bytes, (
+            "Min size not set correctly"
+        )
+        assert min_time == jax.config.jax_persistent_cache_min_compile_time_secs, (
+            "Min time not set correctly"
+        )
         expected_xla = "all" if enable_xla else orig_xla
-        assert (
-            expected_xla == jax.config.jax_persistent_cache_enable_xla_caches
-        ), "XLA caches not set correctly"
+        assert expected_xla == jax.config.jax_persistent_cache_enable_xla_caches, (
+            "XLA caches not set correctly"
+        )
 
     finally:
         if orig_cache_dir is not None: