Terminate humanoid

ss2r/benchmark_suites/mujoco_playground/humanoid/nonepisodic_humanoid.py

@@ -4,6 +4,7 @@
 import jax.numpy as jp
 import mujoco
 from ml_collections import config_dict
+from mujoco import mjx
 from mujoco_playground._src import mjx_env
 from mujoco_playground._src.dm_control_suite import humanoid
 
@@ -21,7 +22,7 @@ def default_config() -> config_dict.ConfigDict:
         episode_length=1000,
         action_repeat=1,
         vision=False,
-        ground_start_probability=1.0,
+        ground_start_probability=0.0,
     )
 
 
@@ -31,13 +32,36 @@ class NonEpisodicHumanoid(humanoid.Humanoid):
     def __init__(
         self,
         move_speed: float = 0.0,
-        config: config_dict.ConfigDict = humanoid.default_config(),
+        config: config_dict.ConfigDict = default_config(),
         config_overrides: Optional[Dict[str, Union[str, int, list[Any]]]] = None,
     ):
         super().__init__(
             move_speed=move_speed, config=config, config_overrides=config_overrides
         )
+
+        self._xml_path = humanoid._XML_PATH.as_posix()
+        mj_spec = mujoco.MjSpec.from_file(
+            filename=str(humanoid._XML_PATH), assets=humanoid.common.get_assets()
+        )
+        self._mj_model = self._modify_model(mj_spec)
+        self._mjx_model = mjx.put_model(self._mj_model)
         self.ground_start_probability = config.ground_start_probability
+        self._post_init()
+
+    def _modify_model(self, mj_spec: mujoco.MjSpec) -> mujoco.MjModel:
+        mj_spec.add_sensor(
+            type=mujoco.mjtSensor.mjSENS_TOUCH,
+            name="head_touch",
+            objtype=mujoco.mjtObj.mjOBJ_SITE,
+            objname="head",
+        )
+        mj_spec.add_sensor(
+            type=mujoco.mjtSensor.mjSENS_TOUCH,
+            name="torso_touch",
+            objtype=mujoco.mjtObj.mjOBJ_SITE,
+            objname="torso",
+        )
+        return mj_spec.compile()
 
     def _post_init(self):
         super()._post_init()
@@ -133,10 +157,27 @@ def normal_init():
 
     def step(self, state: mjx_env.State, action: jax.Array) -> mjx_env.State:
         outs = super().step(state, action)
-        standing = self._head_height(outs.data) > humanoid._STAND_HEIGHT
-        upright = self._torso_upright(outs.data) > 0.9
+        head_height = self._head_height(outs.data)
+        standing = head_height > humanoid._STAND_HEIGHT
+        torso_height = self._torso_upright(outs.data)
+        upright = torso_height > 0.9
         outs.info["cost"] = jp.where(
             standing | upright, jp.zeros_like(outs.reward), jp.ones_like(outs.reward)
         )
-        outs.metrics["reward/on_ground"] = (standing < 0.5).astype(jp.float32)
+        on_ground = (head_height < 0.1) | (torso_height < 0.1)
+        outs.metrics["reward/on_ground"] = on_ground.astype(jp.float32)
+        torso_force = (
+            jp.linalg.norm(
+                mjx_env.get_sensor_data(self.mj_model, state.data, "torso_touch")
+            )
+            > 1000.0
+        )
+        head_force = (
+            jp.linalg.norm(
+                mjx_env.get_sensor_data(self.mj_model, state.data, "head_touch")
+            )
+            > 1000.0
+        )
+        done = on_ground & (torso_force | head_force)
+        outs = outs.replace(done=done.astype(jp.float32))
         return outs

ss2r/benchmark_suites/wrappers.py

@@ -153,12 +153,16 @@ def reset(self, rng):
         state.info["average_reward"] = jp.zeros(rng.shape[:-1])
         metrics = {
             "average_reward": state.info["average_reward"],
+            "alive": jp.ones_like(state.done),
+            "total_steps": state.info["steps"],
         }
         state.info["truncation"] = jp.zeros(rng.shape[:-1])
         state.metrics.update(metrics)
         return state
 
     def step(self, state: State, action: jax.Array) -> State:
+        dead = state.done.astype(jp.bool) | state.info["truncation"].astype(jp.bool)
+
         def f(state, _):
             nstate = self.env.step(state, action)
             maybe_cost = nstate.info.get("cost", None)
@@ -168,13 +172,19 @@ def f(state, _):
         state, (rewards, maybe_costs, maybe_eval_rewards) = jax.lax.scan(
             f, state, (), self.action_repeat
         )
-        sum_rewards = jp.sum(rewards, axis=0)
-        state = state.replace(reward=sum_rewards, done=jp.zeros_like(state.done))
+        sum_rewards = jp.sum(rewards, axis=0) * (1.0 - dead.astype(jp.float32))
+        state = state.replace(reward=sum_rewards)
         if maybe_costs is not None:
-            state.info["cost"] = jp.sum(maybe_costs, axis=0)
+            state.info["cost"] = jp.sum(maybe_costs, axis=0) * (
+                1.0 - dead.astype(jp.float32)
+            )
         if maybe_eval_rewards is not None:
-            state.info["eval_reward"] = jp.sum(maybe_eval_rewards, axis=0)
-        steps = state.info["steps"] + self.action_repeat
+            state.info["eval_reward"] = jp.sum(maybe_eval_rewards, axis=0) * (
+                1.0 - dead
+            )
+        steps = state.info["steps"] + self.action_repeat * (
+            1.0 - dead.astype(jp.float32)
+        )
         sum_rewards /= self.action_repeat
         average_reward = (
             state.info["average_reward"]
@@ -183,6 +193,10 @@ def f(state, _):
         state.info["steps"] = steps
         state.info["average_reward"] = average_reward
         state.metrics["average_reward"] = average_reward
+        # Ignore everything that happens after the first done
+        state.info["truncation"] = dead.astype(jp.float32)
+        state.metrics["alive"] = 1.0 - dead.astype(jp.float32)
+        state.metrics["total_steps"] = steps
         return state
 
 

ss2r/configs/environment/nonepisodic_humanoid.yaml

@@ -8,4 +8,4 @@ task_params:
   episode_length: 300
   sim_dt: 0.0025
   vision: false
-  ground_start_probability: 1.
+  ground_start_probability: 0.