implement zero offset

examples/walking.py

@@ -301,7 +301,7 @@ class WalkingConfig(ksim.PPOConfig):
         help="The acceleration for the angular velocity command, in rad/s^2.",
     )
     angular_velocity_range: tuple[float, float] = xax.field(
-        value=(-0.2, 0.2),
+        value=(0.0, 0.2),
         help="The range for the angular velocity command.",
     )
     angular_velocity_zero_prob: float = xax.field(
@@ -534,7 +534,7 @@ def get_rewards(self, physics_model: ksim.PhysicsModel) -> dict[str, ksim.Reward
                 force_obs="feet_force",
                 ctrl_dt=self.config.ctrl_dt,
                 expected_weight=10.0,
-                scale=-1e-1,
+                scale=1e-1,
             ),
             "motionless_at_rest": ksim.MotionlessAtRestPenalty(scale=1e-2),
         }
@@ -779,6 +779,7 @@ def sample_action(
             # Engine parameters.
             dt=0.004,
             ctrl_dt=0.02,
+            zero_offset_std=math.radians(5.0),
             # Simulation parameters.
             iterations=8,
             ls_iterations=8,

ksim/actuators.py

@@ -19,7 +19,7 @@
 from jaxtyping import Array, PRNGKeyArray, PyTree
 
 from ksim.noise import Noise, NoNoise, RandomVariable, UniformRandomVariable
-from ksim.types import Metadata, PhysicsData, PhysicsModel
+from ksim.types import Metadata, PhysicsModel
 from ksim.utils.mujoco import get_ctrl_data_idx_by_name
 
 logger = logging.getLogger(__name__)
@@ -29,28 +29,32 @@ class Actuators(ABC):
     """Collection of actuators."""
 
     @abstractmethod
-    def get_ctrl(self, action: Array, physics_data: PhysicsData, curriculum_level: Array, rng: PRNGKeyArray) -> Array:
+    def get_ctrl(
+        self,
+        action: Array,
+        qpos: Array,
+        qvel: Array,
+        curriculum_level: Array,
+        rng: PRNGKeyArray,
+    ) -> Array:
         """Get the control signal from the action vector."""
 
-    def get_default_action(self, physics_data: PhysicsData) -> Array:
-        """Get the default action for the actuators."""
-        return physics_data.ctrl
-
 
 class StatefulActuators(Actuators):
     @abstractmethod
     def get_stateful_ctrl(
         self,
         action: Array,
-        physics_data: PhysicsData,
+        qpos: Array,
+        qvel: Array,
         curriculum_level: Array,
         actuator_state: PyTree,
         rng: PRNGKeyArray,
     ) -> tuple[Array, PyTree]:
         """Get the control signal from the action vector."""
 
     @abstractmethod
-    def get_initial_state(self, physics_data: PhysicsData, rng: PRNGKeyArray) -> PyTree:
+    def get_initial_state(self, qpos: Array, qvel: Array, rng: PRNGKeyArray) -> PyTree:
         """Get the initial state for the actuator."""
 
 
@@ -62,7 +66,14 @@ def __init__(self, noise: Noise | None = None) -> None:
 
         self.noise = NoNoise() if noise is None else noise
 
-    def get_ctrl(self, action: Array, physics_data: PhysicsData, curriculum_level: Array, rng: PRNGKeyArray) -> Array:
+    def get_ctrl(
+        self,
+        action: Array,
+        qpos: Array,
+        qvel: Array,
+        curriculum_level: Array,
+        rng: PRNGKeyArray,
+    ) -> Array:
         """Just use the action as the torque, the simplest actuator model."""
         return self.noise.add_noise(action, curriculum_level, rng)
 
@@ -132,13 +143,20 @@ def get_actuator_name(self, joint_name: str) -> str:
         # This can be overridden if necessary.
         return f"{joint_name}_ctrl"
 
-    def get_ctrl(self, action: Array, physics_data: PhysicsData, curriculum_level: Array, rng: PRNGKeyArray) -> Array:
+    def get_ctrl(
+        self,
+        action: Array,
+        qpos: Array,
+        qvel: Array,
+        curriculum_level: Array,
+        rng: PRNGKeyArray,
+    ) -> Array:
         """Get the control signal from the (position) action vector."""
         scaled = action * self.action_scale
 
         pos_rng, tor_rng = jax.random.split(rng)
-        current_pos = physics_data.qpos[7:]  # First 7 are always root pos.
-        current_vel = physics_data.qvel[6:]  # First 6 are always root vel.
+        current_pos = qpos[7:]  # First 7 are always root pos.
+        current_vel = qvel[6:]  # First 6 are always root vel.
 
         # Add position and velocity noise
         target_position = self.action_noise.add_noise(scaled, curriculum_level, pos_rng)
@@ -172,12 +190,19 @@ def __init__(
 
         self.vel_action_noise = NoNoise() if vel_action_noise is None else vel_action_noise
 
-    def get_ctrl(self, action: Array, physics_data: PhysicsData, curriculum_level: Array, rng: PRNGKeyArray) -> Array:
+    def get_ctrl(
+        self,
+        action: Array,
+        qpos: Array,
+        qvel: Array,
+        curriculum_level: Array,
+        rng: PRNGKeyArray,
+    ) -> Array:
         """Get the control signal from the (position and velocity) action vector."""
         pos_rng, vel_rng, tor_rng = jax.random.split(rng, 3)
 
-        current_pos = physics_data.qpos[7:]  # First 7 are always root pos.
-        current_vel = physics_data.qvel[6:]  # First 6 are always root vel.
+        current_pos = qpos[7:]  # First 7 are always root pos.
+        current_vel = qvel[6:]  # First 6 are always root vel.
 
         # Extract position and velocity targets
         target_position = action[: len(current_pos)]
@@ -198,11 +223,6 @@ def get_ctrl(self, action: Array, physics_data: PhysicsData, curriculum_level: A
             self.ctrl_clip,
         )
 
-    def get_default_action(self, physics_data: PhysicsData) -> Array:
-        """Get the default action (zeros) with the correct shape."""
-        qpos_dim = len(physics_data.qpos[7:])
-        return jnp.zeros(qpos_dim * 2)
-
 
 class TorqueBias(TypedDict):
     action: Array
@@ -240,7 +260,8 @@ def __init__(
     def get_stateful_ctrl(
         self,
         action: Array,
-        physics_data: PhysicsData,
+        qpos: Array,
+        qvel: Array,
         curriculum_level: Array,
         actuator_state: TorqueBias,
         rng: PRNGKeyArray,
@@ -252,8 +273,8 @@ def get_stateful_ctrl(
         scaled = action * self.action_scale
 
         pos_rng, tor_rng = jax.random.split(rng)
-        current_pos = physics_data.qpos[7:]  # First 7 are always root pos.
-        current_vel = physics_data.qvel[6:]  # First 6 are always root vel.
+        current_pos = qpos[7:]  # First 7 are always root pos.
+        current_vel = qvel[6:]  # First 6 are always root vel.
 
         # Add position and velocity noise
         target_position = self.action_noise.add_noise(scaled, curriculum_level, pos_rng) + action_bias
@@ -266,8 +287,8 @@ def get_stateful_ctrl(
         ctrl = self.torque_noise.add_noise(ctrl, curriculum_level, tor_rng) + torque_bias
         return jnp.clip(ctrl, -self.ctrl_clip, self.ctrl_clip), actuator_state
 
-    def get_initial_state(self, physics_data: PhysicsData, rng: PRNGKeyArray) -> TorqueBias:
-        shape = physics_data.qpos[..., 7:].shape
+    def get_initial_state(self, qpos: Array, qvel: Array, rng: PRNGKeyArray) -> TorqueBias:
+        shape = qpos[7:].shape
         return {
             "action": self.action_bias.get_random_variable(shape, rng),
             "torque": self.torque_bias.get_random_variable(shape, rng),