[wip] Work towards pipelined sim

Author: amacati

crazyflow/sim/core.py

@@ -14,11 +14,11 @@
 from crazyflow.control.controller import J_INV, Control, Controller, J
 from crazyflow.exception import ConfigError, NotInitializedError
 from crazyflow.sim.fused import (
+    attitude2rpm,
     fused_analytical_dynamics,
     fused_identified_dynamics,
-    fused_masked_attitude2rpm,
-    fused_masked_state2attitude,
     fused_rpms2collective_wrench,
+    state2attitude,
 )
 from crazyflow.sim.integration import Integrator
 from crazyflow.sim.physics import Physics
@@ -125,12 +125,12 @@ def setup_pipeline(self) -> Callable[[int, SimData], SimData]:
         # drones (axis 1) in parallel.
         ctrl_fn = jax.vmap(jax.vmap(self._control_fn()))
         physics_fn = jax.vmap(jax.vmap(self._physics_fn()))
-        integrator_fn = jax.vmap(jax.vmap(self._integrator_fn()))
+        # integrator_fn = jax.vmap(jax.vmap(self._integrator_fn()))
 
         def _step(sim_data: SimData) -> SimData:
             sim_data = ctrl_fn(sim_data)
             sim_data = physics_fn(sim_data)
-            sim_data = integrator_fn(sim_data)
+            # sim_data = integrator_fn(sim_data)
             return sim_data
 
         # ``scan`` can be lowered to a single WhileOp, reducing compilation times while still fusing
@@ -233,7 +233,7 @@ def contacts(self, body: str | None = None) -> Array:
     def _control_fn(self) -> Callable[[SimData], SimData]:
         match self.control:
             case Control.state:
-                return step_state_controller
+                return lambda data: step_attitude_controller(step_state_controller(data))
             case Control.attitude:
                 return step_attitude_controller
             case _:
@@ -242,9 +242,9 @@ def _control_fn(self) -> Callable[[SimData], SimData]:
     def _physics_fn(self) -> Callable[[SimData], SimData]:
         match self.physics:
             case Physics.analytical:
-                return self._step_analytical
+                return analytical_dynamics
             case Physics.sys_id:
-                return self._step_sys_id
+                return identified_dynamics
             case _:
                 raise NotImplementedError(f"Physics mode {self.physics} not implemented")
 
@@ -260,8 +260,7 @@ def _step_sys_id(self):
         # Optional optimization: check if mask.any() before updating the controls. This breaks jax's
         # gradient tracing, so we omit it for now.
         if self.control == Control.state:
-            self.controls = self._masked_state_controls_update(mask, self.controls)
-            self.controls = fused_masked_state2attitude(mask, self.states, self.controls, self.dt)
+            self.controls = state2attitude(mask, self.states, self.controls, self.dt)
         self.controls = self._masked_attitude_controls_update(mask, self.controls)
         self.last_ctrl_steps = self._masked_controls_step_update(
             mask, self.steps, self.last_ctrl_steps
@@ -288,10 +287,9 @@ def _step_analytical(self):
     def _step_emulate_firmware(self) -> SimControls:
         mask = self.controllable
         if self.control == Control.state:
-            self.controls = self._masked_state_controls_update(mask, self.controls)
-            self.controls = fused_masked_state2attitude(mask, self.states, self.controls, self.dt)
+            self.controls = state2attitude(mask, self.states, self.controls, self.dt)
         self.controls = self._masked_attitude_controls_update(mask, self.controls)
-        return fused_masked_attitude2rpm(mask, self.states, self.controls, self.dt)
+        return attitude2rpm(mask, self.states, self.controls, self.dt)
 
     @staticmethod
     def _sync_mjx(states: SimState, mjx_data: Data, mjx_model: Model) -> Data:
@@ -323,64 +321,52 @@ def _sync_mjx_full(states: SimState, mjx_data: Data, mjx_model: Model) -> Data:
     @staticmethod
     @jax.jit
     def _masked_states_reset(mask: Array, states: SimState, defaults: SimState) -> SimState:
-        mask_3d = mask[:, None, None]
-        states = states.replace(pos=jnp.where(mask_3d, defaults.pos, states.pos))
-        states = states.replace(quat=jnp.where(mask_3d, defaults.quat, states.quat))
-        states = states.replace(vel=jnp.where(mask_3d, defaults.vel, states.vel))
-        states = states.replace(ang_vel=jnp.where(mask_3d, defaults.ang_vel, states.ang_vel))
-        states = states.replace(rpy_rates=jnp.where(mask_3d, defaults.rpy_rates, states.rpy_rates))
-        return states
+        mask = mask.reshape(-1, 1, 1)
+        return jax.tree.map(lambda x, y: jnp.where(mask, y, x), states, defaults)
 
     @staticmethod
     @jax.jit
     def _masked_controls_reset(
         mask: Array, controls: SimControls, defaults: SimControls
     ) -> SimControls:
-        mask = mask[:, None, None]
-        controls = controls.replace(
-            state=jnp.where(mask, defaults.state, controls.state),
-            attitude=jnp.where(mask, defaults.attitude, controls.attitude),
-            thrust=jnp.where(mask, defaults.thrust, controls.thrust),
-            rpms=jnp.where(mask, defaults.rpms, controls.rpms),
-            rpy_err_i=jnp.where(mask, defaults.rpy_err_i, controls.rpy_err_i),
-            pos_err_i=jnp.where(mask, defaults.pos_err_i, controls.pos_err_i),
-            last_rpy=jnp.where(mask, defaults.last_rpy, controls.last_rpy),
-            staged_attitude=jnp.where(mask, defaults.staged_attitude, controls.staged_attitude),
-            staged_state=jnp.where(mask, defaults.staged_state, controls.staged_state),
-        )
-        return controls
+        mask = mask.reshape(-1, 1, 1)
+        return jax.tree.map(lambda x, y: jnp.where(mask, y, x), controls, defaults)
 
     @staticmethod
     @jax.jit
     def _masked_params_reset(mask: Array, params: SimParams, defaults: SimParams) -> SimParams:
-        params = params.replace(mass=jnp.where(mask[:, None, None], defaults.mass, params.mass))
-        mask_4d = mask[:, None, None, None]
-        params = params.replace(J=jnp.where(mask_4d, defaults.J, params.J))
-        params = params.replace(J_INV=jnp.where(mask_4d, defaults.J_INV, params.J_INV))
+        mask = mask.reshape(-1, 1, 1)
+        params = params.replace(mass=jnp.where(mask, defaults.mass, params.mass))
+        # J and J_INV are matrices -> we need (W, D, N, N) = 4 dims
+        mask = mask.reshape(-1, 1, 1, 1)
+        params = params.replace(J=jnp.where(mask, defaults.J, params.J))
+        params = params.replace(J_INV=jnp.where(mask, defaults.J_INV, params.J_INV))
         return params
 
     @staticmethod
     @partial(jax.jit, static_argnames="device")
     def _attitude_control(cmd: Array, controls: SimControls, device: str) -> SimControls:
+        """Stage the desired attitude for all drones in all worlds.
+
+        We need to stage the attitude commands because the sys_id physics mode operates directly on
+        the attitude command. 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 commands, we ensure that the physics module sees the old commands until the
+        controller updates at its correct frequency.
+        """
         return controls.replace(staged_attitude=jnp.array(cmd, device=device))
 
     @staticmethod
     @partial(jax.jit, static_argnames="device")
     def _state_control(cmd: Array, controls: SimControls, device: str) -> SimControls:
-        return controls.replace(staged_state=jnp.array(cmd, device=device))
+        return controls.replace(state=jnp.array(cmd, device=device))
 
     @staticmethod
     @jax.jit
     def _masked_attitude_controls_update(mask: Array, controls: SimControls) -> SimControls:
         cmd, staged_cmd = controls.attitude, controls.staged_attitude
         return controls.replace(attitude=jnp.where(mask[:, None, None], staged_cmd, cmd))
 
-    @staticmethod
-    @jax.jit
-    def _masked_state_controls_update(mask: Array, controls: SimControls) -> SimControls:
-        cmd, staged_cmd = controls.state, controls.staged_state
-        return controls.replace(state=jnp.where(mask[:, None, None], staged_cmd, cmd))
-
     @staticmethod
     @jax.jit
     def _masked_controls_step_update(mask: Array, steps: Array, last_ctrl_steps: Array) -> Array:
@@ -403,24 +389,30 @@ def contacts(geom_start: int, geom_count: int, data: Data) -> Array:
 
 
 def step_state_controller(data: SimData) -> SimData:
+    """Compute the updated controls for the state controller."""
     controls = data.controls
-    mask = controllable(data.steps, controls.steps, data.freq, controls.state_freq)
-    controls = commit_state_controls(mask, controls)
+    mask = controllable(data.steps, controls.state_steps, data.freq, controls.state_freq)
+    controls = controls.replace(state_steps=jnp.where(mask, data.steps, controls.state_steps))
     controls = state2attitude(mask, data.states, controls, 1 / data.freq)
     return data.replace(controls=controls)
 
 
-def controllable(step: Array, ctrl_step: Array, ctrl_freq: int, freq: int) -> Array:
-    return (step - ctrl_step) >= (freq / ctrl_freq)
+def step_attitude_controller(data: SimData) -> SimData:
+    """Compute the updated controls for the attitude controller."""
+    controls = data.controls
+    mask = controllable(data.steps, controls.attitude_steps, data.freq, controls.attitude_freq)
+    controls = commit_attitude_controls(mask, controls)
+    controls = attitude2rpm(mask, data.states, controls, 1 / data.freq)
+    return data.replace(controls=controls)
 
 
-def commit_state_controls(mask: Array, controls: SimControls) -> SimControls:
-    cmd, staged_cmd = controls.state, controls.staged_state
-    return controls.replace(state=jnp.where(mask[:, None, None], staged_cmd, cmd))
+def controllable(step: Array, ctrl_step: Array, ctrl_freq: int, freq: int) -> Array:
+    return (step - ctrl_step) >= (freq / ctrl_freq)
 
 
-def step_attitude_controller(data: SimData) -> SimData:
-    pass
+def commit_attitude_controls(mask: Array, controls: SimControls) -> SimControls:
+    cmd, staged_cmd = controls.attitude, controls.staged_attitude
+    return controls.replace(attitude=jnp.where(mask.reshape(-1, 1, 1), staged_cmd, cmd))
 
 
 mjx_kinematics = jax.vmap(mjx.kinematics, in_axes=(None, 0))

crazyflow/sim/fused.py

@@ -11,7 +11,8 @@
 from jax import Array
 from jax.scipy.spatial.transform import Rotation as R
 
-from crazyflow.control.controller import attitude2rpm, state2attitude
+from crazyflow.control.controller import attitude2rpm as attitude2rpm_ctrl
+from crazyflow.control.controller import state2attitude as state2attitude_ctrl
 from crazyflow.sim.physics import analytical_dynamics, identified_dynamics, rpms2collective_wrench
 from crazyflow.sim.structs import SimControls, SimParams, SimState
 
@@ -66,9 +67,7 @@ def fused_analytical_dynamics(
 @jax.jit
 @partial(jax.vmap, in_axes=(0, 0, 0, None))
 @partial(jax.vmap, in_axes=(None, 0, 0, None))
-def fused_masked_state2attitude(
-    mask: Array, state: SimState, cmd: SimControls, dt: float
-) -> SimControls:
+def state2attitude(mask: Array, state: SimState, cmd: SimControls, dt: float) -> SimControls:
     """Compute the next desired collective thrust and roll/pitch/yaw of the drone.
 
     Note:
@@ -86,7 +85,7 @@ def fused_masked_state2attitude(
         dt: The simulation time step.
     """
     des_pos, des_vel, des_yaw = cmd.state[:3], cmd.state[3:6], cmd.state[9].reshape((1,))
-    attitude, pos_err_i = state2attitude(
+    attitude, pos_err_i = state2attitude_ctrl(
         state.pos, state.vel, state.quat, des_pos, des_vel, des_yaw, cmd.pos_err_i, dt
     )
     # Non-branching selection depending on the mask. XLA should be able to optimize a short path
@@ -99,9 +98,7 @@ def fused_masked_state2attitude(
 @jax.jit
 @partial(jax.vmap, in_axes=(0, 0, 0, None))
 @partial(jax.vmap, in_axes=(None, 0, 0, None))
-def fused_masked_attitude2rpm(
-    mask: Array, state: SimState, cmd: SimControls, dt: float
-) -> SimControls:
+def attitude2rpm(mask: Array, state: SimState, cmd: SimControls, dt: float) -> SimControls:
     """Compute the next desired RPMs of the drone.
 
     Note:
@@ -114,7 +111,7 @@ def fused_masked_attitude2rpm(
         cmd: The current simulation controls.
         dt: The simulation time step.
     """
-    rpms, rpy_err_i = attitude2rpm(cmd.attitude, state.quat, cmd.last_rpy, cmd.rpy_err_i, dt)
+    rpms, rpy_err_i = attitude2rpm_ctrl(cmd.attitude, state.quat, cmd.last_rpy, cmd.rpy_err_i, dt)
     # Non-branching selection depending on the mask. See fused_masked_state2attitude for more info.
     rpms = jnp.where(mask, rpms, cmd.rpms)
     rpy_err_i = jnp.where(mask, rpy_err_i, cmd.rpy_err_i)

crazyflow/sim/integration.py

@@ -1,7 +1,57 @@
 from enum import Enum
+from typing import Callable
+
+from jax import Array
+from scipy.spatial.transform import Rotation as R
+
+from crazyflow.sim.structs import SimControls
 
 
 class Integrator(str, Enum):
     euler = "euler"
     # TODO: Implement rk4
     default = euler  # TODO: Replace with rk4
+
+
+def euler(
+    deriv_fn: Callable[[Array, Array, Array, Array, SimControls, float], tuple[Array, Array]],
+    pos: Array,
+    quat: Array,
+    vel: Array,
+    rpy_rates: Array,
+    control: SimControls,
+    dt: float,
+) -> tuple[Array, Array, Array, Array]:
+    acc, rpy_rates_deriv = deriv_fn(pos, quat, vel, rpy_rates, control, dt)
+    return _apply(pos, quat, vel, rpy_rates, acc, rpy_rates_deriv, dt)
+
+
+def rk4(
+    deriv_fn: Callable[[Array, Array, Array, Array, SimControls, float], tuple[Array, Array]],
+    pos: Array,
+    quat: Array,
+    vel: Array,
+    rpy_rates: Array,
+    control: SimControls,
+    dt: float,
+) -> tuple[Array, Array, Array, Array]:
+    raise NotImplementedError("RK4 not implemented")
+
+
+def _apply(
+    pos: Array,
+    quat: Array,
+    vel: Array,
+    rpy_rates: Array,
+    acc: Array,
+    rpy_rates_deriv: Array,
+    dt: float,
+) -> tuple[Array, Array, Array, Array]:
+    rot = R.from_quat(quat)
+    rpy_rates_local = rot.apply(rpy_rates, inverse=True)
+    rpy_rates_deriv_local = rot.apply(rpy_rates_deriv, inverse=True)
+    next_pos = pos + vel * dt
+    next_vel = vel + acc * dt
+    next_rot = R.from_euler("xyz", rot.as_euler("xyz") + rpy_rates_local * dt)
+    next_rpy_rates_local = rpy_rates_local + rpy_rates_deriv_local * dt
+    return next_pos, next_rot.as_quat(), next_vel, next_rot.apply(next_rpy_rates_local)

crazyflow/sim/physics.py

@@ -6,6 +6,7 @@
 
 from crazyflow.constants import ARM_LEN, GRAVITY, SIGN_MIX_MATRIX
 from crazyflow.control.controller import KF, KM
+from crazyflow.sim.structs import SimControls, SimParams, SimState
 
 SYS_ID_PARAMS = {
     "acc_k1": 20.91,
@@ -69,6 +70,30 @@ def identified_dynamics(
     return next_pos, next_quat, next_vel, next_rpy_rates
 
 
+def identified_dynamics_dx(
+    state: SimState, controls: SimControls, dt: float
+) -> tuple[Array, Array]:
+    """Derivative of the identified dynamics state."""
+    collective_thrust, attitude = controls.attitude[0], controls.attitude[1:]
+    rot = R.from_quat(state.quat)
+    thrust = rot.apply(jnp.array([0, 0, collective_thrust]))
+    drift = rot.apply(jnp.array([0, 0, 1]))
+    a1, a2 = SYS_ID_PARAMS["acc_k1"], SYS_ID_PARAMS["acc_k2"]
+    acc = thrust * a1 + drift * a2 - jnp.array([0, 0, GRAVITY])
+    # rpy_rates_deriv have no real meaning in this context, since the identified dynamics set the
+    # rpy_rates to the commanded values directly. However, since we use a unified integration
+    # interface for all physics models, we cannot access states directly. Instead, we calculate
+    # which rpy_rates_deriv would have resulted in the desired rpy_rates, and return that.
+    roll_cmd, pitch_cmd, yaw_cmd = attitude
+    rpy = rot.as_euler("xyz")
+    roll_rate = SYS_ID_PARAMS["roll_alpha"] * rpy[0] + SYS_ID_PARAMS["roll_beta"] * roll_cmd
+    pitch_rate = SYS_ID_PARAMS["pitch_alpha"] * rpy[1] + SYS_ID_PARAMS["pitch_beta"] * pitch_cmd
+    yaw_rate = SYS_ID_PARAMS["yaw_alpha"] * rpy[2] + SYS_ID_PARAMS["yaw_beta"] * yaw_cmd
+    rpy_rates = jnp.array([roll_rate, pitch_rate, yaw_rate])
+    rpy_rates_deriv = (rpy_rates - rot.apply(state.rpy_rates, inverse=True)) / dt
+    return acc, rpy_rates_deriv
+
+
 def analytical_dynamics(
     forces: Array,
     torques: Array,
@@ -91,14 +116,25 @@ def analytical_dynamics(
     # Update state.
     next_pos = pos + vel * dt
     next_vel = vel + acc * dt
-    next_rot = R.from_euler("xyz", R.from_quat(quat).as_euler("xyz") + rpy_rates_local * dt)
+    next_rot = R.from_euler("xyz", rot.as_euler("xyz") + rpy_rates_local * dt)
     next_quat = next_rot.as_quat()
     # Convert rpy rates back to global frame
     next_rpy_rates_local = rpy_rates_local + rpy_rates_deriv_local * dt
     next_rpy_rates = next_rot.apply(next_rpy_rates_local)  # Always give rpy rates in world frame
     return next_pos, next_quat, next_vel, next_rpy_rates
 
 
+def analytical_dynamics_dx(
+    forces: Array, torques: Array, state: SimState, params: SimParams
+) -> tuple[Array, Array]:
+    """Derivative of the analytical dynamics state."""
+    rot = R.from_quat(state.quat)
+    torques_local = rot.apply(torques, inverse=True)
+    acc = forces / params.mass - jnp.array([0, 0, GRAVITY])
+    rpy_rates_deriv = rot.apply(params.J_INV @ torques_local)
+    return acc, rpy_rates_deriv
+
+
 def rpms2collective_wrench(
     rpms: Array, quat: Array, rpy_rates: Array, J: Array
 ) -> tuple[Array, Array]:

crazyflow/sim/structs.py

@@ -26,7 +26,6 @@ def default_state(n_worlds: int, n_drones: int, device: Device) -> SimState:
 @dataclass
 class SimControls:
     state: Array  # (N, M, 13)
-    staged_state: Array  # (N, M, 13)
     attitude: Array  # (N, M, 4)
     staged_attitude: Array  # (N, M, 4)
     thrust: Array  # (N, M, 4)
@@ -40,7 +39,6 @@ def default_controls(n_worlds: int, n_drones: int, device: Device) -> SimControl
     """Create a default set of controls for the simulation."""
     return SimControls(
         state=jnp.zeros((n_worlds, n_drones, 13), device=device),
-        staged_state=jnp.zeros((n_worlds, n_drones, 13), device=device),
         attitude=jnp.zeros((n_worlds, n_drones, 4), device=device),
         staged_attitude=jnp.zeros((n_worlds, n_drones, 4), device=device),
         thrust=jnp.zeros((n_worlds, n_drones, 4), device=device),