implement zero offset (#538)

* implement zero offset

* thought for a bit

* remove biased actuator model

* remove biased thing

* fix shape mismatch

* lint

Author: codekansas

examples/kbot/train.py

@@ -452,11 +452,9 @@ def get_actuators(
         metadata: ksim.Metadata | None = None,
     ) -> ksim.Actuators:
         assert metadata is not None, "Metadata is required"
-        return ksim.BiasedPositionActuators(
+        return ksim.PositionActuators(
             physics_model=physics_model,
             metadata=metadata,
-            action_bias=ksim.UniformRandomVariable(mean=0.0, mag=math.radians(3.0)),
-            torque_bias=ksim.UniformRandomVariable(mean=0.0, mag=3.0),
             action_noise=ksim.AdditiveGaussianNoise(std=0.01),
             torque_noise=ksim.AdditiveGaussianNoise(std=0.01),
         )
@@ -891,6 +889,8 @@ def sample_action(
             ctrl_dt=0.02,
             action_latency_range=(0.001, 0.01),  # Simulate 1-10ms of latency.
             drop_action_prob=0.05,  # Drop 5% of commands.
+            zero_offset_std=math.radians(3.0),
+            zero_offset_mag=math.radians(3.0),
             # Visualization parameters.
             # If running this on Mac and you are getting segfaults,
             # you might need to disable `render_markers`

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,8 @@ def sample_action(
             # Engine parameters.
             dt=0.004,
             ctrl_dt=0.02,
+            zero_offset_std=math.radians(3.0),
+            zero_offset_mag=math.radians(3.0),
             # Simulation parameters.
             iterations=8,
             ls_iterations=8,

ksim/actuators.py

@@ -6,20 +6,18 @@
     "TorqueActuators",
     "PositionActuators",
     "PositionVelocityActuator",
-    "BiasedPositionActuators",
 ]
 
 import logging
 from abc import ABC, abstractmethod
-from typing import TypedDict
 
 import chex
 import jax
 import jax.numpy as jnp
 from jaxtyping import Array, PRNGKeyArray, PyTree
 
-from ksim.noise import Noise, NoNoise, RandomVariable, UniformRandomVariable
-from ksim.types import Metadata, PhysicsData, PhysicsModel
+from ksim.noise import Noise, NoNoise
+from ksim.types import Metadata, PhysicsModel
 from ksim.utils.mujoco import get_ctrl_data_idx_by_name
 
 logger = logging.getLogger(__name__)
@@ -29,28 +27,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 +64,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 +141,22 @@ 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.
+
+        # Calling function removes root position and velocity.
+        current_pos = qpos
+        current_vel = qvel
 
         # Add position and velocity noise
         target_position = self.action_noise.add_noise(scaled, curriculum_level, pos_rng)
@@ -172,12 +190,20 @@ 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.
+        # Calling function removes root position and velocity.
+        current_pos = qpos
+        current_vel = qvel
 
         # Extract position and velocity targets
         target_position = action[: len(current_pos)]
@@ -197,78 +223,3 @@ def get_ctrl(self, action: Array, physics_data: PhysicsData, curriculum_level: A
             -self.ctrl_clip,
             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
-    torque: Array
-
-
-class BiasedPositionActuators(PositionActuators, StatefulActuators):
-    """Adds some random bias to the position action to simulate imperfect actuators."""
-
-    def __init__(
-        self,
-        physics_model: PhysicsModel,
-        metadata: Metadata,
-        action_bias: RandomVariable | float,
-        torque_bias: RandomVariable | float,
-        action_noise: Noise | None = None,
-        torque_noise: Noise | None = None,
-        action_scale: float = 1.0,
-    ) -> None:
-        super().__init__(
-            physics_model=physics_model,
-            metadata=metadata,
-            action_noise=action_noise,
-            torque_noise=torque_noise,
-            action_scale=action_scale,
-        )
-
-        if not isinstance(action_bias, RandomVariable):
-            action_bias = UniformRandomVariable(mean=0.0, mag=action_bias)
-        if not isinstance(torque_bias, RandomVariable):
-            torque_bias = UniformRandomVariable(mean=0.0, mag=torque_bias)
-        self.action_bias = action_bias
-        self.torque_bias = torque_bias
-
-    def get_stateful_ctrl(
-        self,
-        action: Array,
-        physics_data: PhysicsData,
-        curriculum_level: Array,
-        actuator_state: TorqueBias,
-        rng: PRNGKeyArray,
-    ) -> tuple[Array, TorqueBias]:
-        """Get the control signal from the (position) action vector."""
-        action_bias = actuator_state["action"]
-        torque_bias = actuator_state["torque"]
-
-        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.
-
-        # Add position and velocity noise
-        target_position = self.action_noise.add_noise(scaled, curriculum_level, pos_rng) + action_bias
-        target_velocity = jnp.zeros_like(action)
-
-        pos_delta = target_position - current_pos
-        vel_delta = target_velocity - current_vel
-
-        ctrl = self.kps * pos_delta + self.kds * vel_delta
-        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
-        return {
-            "action": self.action_bias.get_random_variable(shape, rng),
-            "torque": self.torque_bias.get_random_variable(shape, rng),
-        }