prepare reference motion for amp (#63)

* g1 obs noise scales

* n

* update randomizers

* setup works again

* fixed elbows

* save debugging setup

* fix noise, change back to accel

* more noise

* clean qpos logic

* add orientation logic

* add terrain setup

* push updated rnn

* format

* run env works again

* training works as well

* format

* add back gyro

* wip, remove old reference logic

* mypy lint

* update jax

* debug velocity rewrad

* proper input shapes

* add rnn version

* add rnn setup

* save functional setup

* add literal typin

---------

Co-authored-by: Wesley Maa <[email protected]>

Author: budzianowski

ksim_kbot/common.py

@@ -19,6 +19,10 @@
     slice_update,
     update_data_field,
 )
+from ksim.utils.priors import (
+    MotionReferenceData,
+    get_local_xpos,
+)
 from mujoco import mjx
 
 
@@ -491,3 +495,47 @@ def _apply_random_angular_velocity_push(
     def get_initial_event_state(self, rng: PRNGKeyArray) -> Array:
         minval, maxval = self.interval_range
         return jax.random.uniform(rng, (), minval=minval, maxval=maxval)
+
+
+@attrs.define(frozen=True, kw_only=True)
+class ReferenceQposObservation(ksim.Observation):
+    """Observation for the reference joint positions."""
+
+    reference_motion_data: MotionReferenceData
+    speed: float = attrs.field(default=1.0)
+
+    def observe(self, state: ksim.ObservationInput, curriculum_level: Array, rng: PRNGKeyArray) -> Array:
+        physics_state = state.physics_state
+        effective_time = physics_state.data.time * self.speed
+        reference_qpos_at_time = self.reference_motion_data.get_qpos_at_time(effective_time)
+        return reference_qpos_at_time[..., 7:]
+
+
+@attrs.define(frozen=True, kw_only=True)
+class ReferenceLocalXposObservation(ksim.Observation):
+    """Observation for the reference local cartesian positions of tracked bodies."""
+
+    reference_motion_data: MotionReferenceData
+    tracked_body_ids: tuple[int, ...]
+
+    def observe(self, state: ksim.ObservationInput, curriculum_level: Array, rng: PRNGKeyArray) -> Array:
+        physics_state = state.physics_state
+        target_pos_dict = self.reference_motion_data.get_cartesian_pose_at_time(physics_state.data.time)
+        target_pos_list = [target_pos_dict[body_id] for body_id in self.tracked_body_ids]
+        return jnp.concatenate(target_pos_list, axis=-1)
+
+
+@attrs.define(frozen=True, kw_only=True)
+class TrackedLocalXposObservation(ksim.Observation):
+    """Observation for the current local cartesian positions of tracked bodies."""
+
+    tracked_body_ids: tuple[int, ...]
+    mj_base_id: int
+
+    def observe(self, state: ksim.ObservationInput, curriculum_level: Array, rng: PRNGKeyArray) -> Array:
+        physics_state = state.physics_state
+        tracked_positions_list: list[Array] = []
+        for body_id in self.tracked_body_ids:
+            body_pos = get_local_xpos(physics_state.data.xpos, body_id, self.mj_base_id)
+            tracked_positions_list.append(jnp.array(body_pos))
+        return jnp.concatenate(tracked_positions_list, axis=-1)

ksim_kbot/rewards.py

@@ -3,7 +3,7 @@
 If some logic will become more general, we can move it to ksim or xax.
 """
 
-from typing import Self
+from typing import Literal, Self
 
 import attrs
 import jax.numpy as jnp
@@ -434,3 +434,40 @@ def gait_phase(
         stance = xax.cubic_bezier_interpolation(jnp.array(0), swing_height, 2 * x)
         swing = xax.cubic_bezier_interpolation(swing_height, jnp.array(0), 2 * x - 1)
         return jnp.where(x <= 0.5, stance, swing)
+
+
+@attrs.define(frozen=True)
+class TargetLinearVelocityReward(ksim.Reward):
+    """Reward for forward motion."""
+
+    index: Literal["x", "y", "z"] = attrs.field(default="x")
+    target_vel: float = attrs.field(default=0.0)
+    norm: xax.NormType = attrs.field(default="l1")
+    monotonic_fn: Literal["exp", "inv"] = attrs.field(default="inv")
+    temp: float = attrs.field(default=1.0)
+
+    def get_reward(self, trajectory: ksim.Trajectory) -> Array:
+        vel = trajectory.qvel[..., ksim.cartesian_index_to_dim(self.index)]
+        error = xax.get_norm(vel - self.target_vel, self.norm)
+        return ksim.norm_to_reward(error, temp=self.temp, monotonic_fn=self.monotonic_fn)
+
+    def get_name(self) -> str:
+        return f"{self.index}_{super().get_name()}"
+
+
+@attrs.define(frozen=True, kw_only=True)
+class TargetHeightReward(ksim.Reward):
+    """Reward for reaching a target height."""
+
+    target_height: float = attrs.field(default=1.0)
+    norm: xax.NormType = attrs.field(default="l1")
+    temp: float = attrs.field(default=1.0)
+    monotonic_fn: Literal["exp", "inv"] = attrs.field(default="inv")
+
+    def get_reward(self, trajectory: ksim.Trajectory) -> Array:
+        qpos = trajectory.qpos
+        error = qpos[..., 2] - self.target_height
+        reward_value = ksim.norm_to_reward(
+            xax.get_norm(error, self.norm), temp=self.temp, monotonic_fn=self.monotonic_fn
+        )
+        return reward_value

ksim_kbot/walking/walking.py

@@ -530,7 +530,7 @@ def sample_action(
     # To run training, use the following command:
     #   python -m ksim_kbot.walking.walking
     # To visualize the environment, use the following command:
-    #   python -m ksim_kbot.walking.walking run_environment=True
+    #   python -m ksim_kbot.walking.walking run_model_viewer=True
     # On MacOS or other devices with less memory, you can change the number
     # of environments and batch size to reduce memory usage. Here's an example
     # from the command line:

ksim_kbot/walking/walking_joystick.py

@@ -65,13 +65,14 @@ def __init__(
         self,
         key: PRNGKeyArray,
         *,
+        num_inputs: int,
         min_std: float,
         max_std: float,
         var_scale: float,
         mean_scale: float,
     ) -> None:
         self.mlp = eqx.nn.MLP(
-            in_size=NUM_INPUTS,
+            in_size=num_inputs,
             out_size=NUM_OUTPUTS * 2,
             width_size=256,
             depth=5,
@@ -133,9 +134,9 @@ class KbotCritic(eqx.Module):
 
     mlp: eqx.nn.MLP
 
-    def __init__(self, key: PRNGKeyArray) -> None:
+    def __init__(self, key: PRNGKeyArray, *, num_inputs: int) -> None:
         self.mlp = eqx.nn.MLP(
-            in_size=NUM_CRITIC_INPUTS,
+            in_size=num_inputs,
             out_size=1,  # Always output a single critic value.
             width_size=256,
             depth=5,
@@ -193,16 +194,30 @@ def forward(
 class KbotModel(eqx.Module):
     actor: KbotActor
     critic: KbotCritic
+    num_inputs: int = eqx.static_field()
+    num_critic_inputs: int = eqx.static_field()
 
-    def __init__(self, key: PRNGKeyArray) -> None:
+    def __init__(
+        self,
+        key: PRNGKeyArray,
+        *,
+        num_inputs: int,
+        num_critic_inputs: int,
+    ) -> None:
+        self.num_inputs = num_inputs
+        self.num_critic_inputs = num_critic_inputs
         self.actor = KbotActor(
             key,
+            num_inputs=num_inputs,
             min_std=0.01,
             max_std=1.0,
             var_scale=1.0,
             mean_scale=1.0,
         )
-        self.critic = KbotCritic(key)
+        self.critic = KbotCritic(
+            key,
+            num_inputs=num_critic_inputs,
+        )
 
 
 @dataclass
@@ -559,7 +574,11 @@ def get_terminations(self, physics_model: ksim.PhysicsModel) -> list[ksim.Termin
         return [common.GVecTermination.create(physics_model, sensor_name="upvector_origin")]
 
     def get_model(self, key: PRNGKeyArray) -> KbotModel:
-        return KbotModel(key)
+        return KbotModel(
+            key,
+            num_inputs=NUM_INPUTS,
+            num_critic_inputs=NUM_CRITIC_INPUTS,
+        )
 
     def get_initial_carry(self, rng: PRNGKeyArray) -> tuple[Array, Array]:
         return None, None
@@ -724,11 +743,13 @@ def on_after_checkpoint_save(self, ckpt_path: Path, state: xax.State) -> xax.Sta
             if self.config.only_save_most_recent
             else ckpt_path.parent / f"tf_model_{state.num_steps}"
         )
+
         export(
             model_fn,
             input_shapes,  # type: ignore [arg-type]
             tf_path,
         )
+
         return state
 
 
@@ -737,7 +758,7 @@ def on_after_checkpoint_save(self, ckpt_path: Path, state: xax.State) -> xax.Sta
     # To run training, use the following command:
     # python -m ksim_kbot.walking.walking_joystick disable_multiprocessing=True
     # To visualize the environment, use the following command:
-    # python -m ksim_kbot.walking.walking_joystick run_environment=True \
+    # python -m ksim_kbot.walking.walking_joystick run_model_viewer=True \
     #  run_environment_num_seconds=1 \
     #  run_environment_save_path=videos/test.mp4
     KbotWalkingTask.launch(
@@ -772,6 +793,5 @@ def on_after_checkpoint_save(self, ckpt_path: Path, state: xax.State) -> xax.Sta
             gait_freq_upper=1.5,
             reward_clip_min=0.0,
             reward_clip_max=1000.0,
-            stand_still_threshold=0.0,  # no stand still reward
         ),
     )

ksim_kbot/walking/walking_joystick_rnn.py

@@ -14,6 +14,7 @@
 import mujoco
 import xax
 from jaxtyping import Array, PRNGKeyArray
+from mujoco import mjx
 from mujoco_scenes.mjcf import load_mjmodel
 from xax.nn.export import export
 
@@ -300,11 +301,11 @@ def get_mujoco_model(self) -> mujoco.MjModel:
         mj_model = load_mjmodel(mjcf_path, scene=self.config.terrain_type)
 
         # NOTE: test the difference
-        # mj_model.opt.timestep = jnp.array(self.config.dt)
-        # mj_model.opt.iterations = 6
-        # mj_model.opt.ls_iterations = 6
-        # mj_model.opt.disableflags = mjx.DisableBit.EULERDAMP
-        # mj_model.opt.solver = mjx.SolverType.CG
+        mj_model.opt.timestep = jnp.array(self.config.dt)
+        mj_model.opt.iterations = 6
+        mj_model.opt.ls_iterations = 6
+        mj_model.opt.disableflags = mjx.DisableBit.EULERDAMP
+        mj_model.opt.solver = mjx.SolverType.CG
 
         return mj_model
 
@@ -534,7 +535,7 @@ def on_after_checkpoint_save(self, ckpt_path: Path, state: xax.State) -> xax.Sta
     # To run training, use the following command:
     #   python -m ksim_kbot.walking.walking_joystick_rnn
     # To visualize the environment, use the following command:
-    #   python -m ksim_kbot.walking.walking_joystick_rnn run_environment=True
+    #   python -m ksim_kbot.walking.walking_joystick_rnn run_model_viewer=True
     KbotWalkingJoystickRNNTask.launch(
         KbotWalkingJoystickRNNTaskConfig(
             num_envs=4096,