Replace hardcoded observation size with jax.eval_shape.

PiperOrigin-RevId: 720599097
Change-Id: I71b7914417b389155c716b08c08cf5781adfdfe5

mujoco_playground/_src/dm_control_suite/acrobot.py

@@ -142,10 +142,6 @@ def xml_path(self) -> str:
   def action_size(self) -> int:
     return self.mjx_model.nu
 
-  @property
-  def observation_size(self) -> mjx_env.ObservationSize:
-    return 6
-
   @property
   def mj_model(self) -> mujoco.MjModel:
     return self._mj_model

mujoco_playground/_src/dm_control_suite/ball_in_cup.py

@@ -116,10 +116,6 @@ def xml_path(self) -> str:
   def action_size(self) -> int:
     return self.mjx_model.nu
 
-  @property
-  def observation_size(self) -> mjx_env.ObservationSize:
-    return 8
-
   @property
   def mj_model(self) -> mujoco.MjModel:
     return self._mj_model

mujoco_playground/_src/dm_control_suite/cartpole.py

@@ -278,10 +278,6 @@ def xml_path(self) -> str:
   def action_size(self) -> int:
     return self.mjx_model.nu
 
-  @property
-  def observation_size(self) -> mjx_env.ObservationSize:
-    return 5
-
   @property
   def mj_model(self) -> mujoco.MjModel:
     return self._mj_model

mujoco_playground/_src/dm_control_suite/cheetah.py

@@ -135,10 +135,6 @@ def xml_path(self) -> str:
   def action_size(self) -> int:
     return self.mjx_model.nu
 
-  @property
-  def observation_size(self) -> mjx_env.ObservationSize:
-    return 17
-
   @property
   def mj_model(self) -> mujoco.MjModel:
     return self._mj_model

mujoco_playground/_src/dm_control_suite/dm_control_suite_test.py

@@ -33,7 +33,7 @@ def test_can_create_all_environments(self, env_name: str) -> None:
     state = jax.jit(env.reset)(jax.random.PRNGKey(42))
     state = jax.jit(env.step)(state, jp.zeros(env.action_size))
     self.assertIsNotNone(state)
-    self.assertEqual(state.obs.shape[0], env.observation_size)
+    self.assertEqual(state.obs.shape, env.observation_size)
     self.assertFalse(jp.isnan(state.data.qpos).any())
 
 

mujoco_playground/_src/dm_control_suite/finger.py

@@ -188,10 +188,6 @@ def sim_dt(self) -> float:
   def action_size(self) -> int:
     return self.mjx_model.nu
 
-  @property
-  def observation_size(self) -> mjx_env.ObservationSize:
-    return 9
-
   @property
   def mj_model(self) -> mujoco.MjModel:
     return self._mj_model
@@ -344,10 +340,6 @@ def xml_path(self) -> str:
   def action_size(self) -> int:
     return self.mjx_model.nu
 
-  @property
-  def observation_size(self) -> mjx_env.ObservationSize:
-    return 12
-
   @property
   def mj_model(self) -> mujoco.MjModel:
     return self._mj_model

mujoco_playground/_src/dm_control_suite/fish.py

@@ -187,10 +187,6 @@ def xml_path(self) -> str:
   def action_size(self) -> int:
     return self.mjx_model.nu
 
-  @property
-  def observation_size(self) -> mjx_env.ObservationSize:
-    return 24
-
   @property
   def mj_model(self) -> mujoco.MjModel:
     return self._mj_model

mujoco_playground/_src/dm_control_suite/hopper.py

@@ -194,10 +194,6 @@ def xml_path(self) -> str:
   def action_size(self) -> int:
     return self.mjx_model.nu
 
-  @property
-  def observation_size(self) -> mjx_env.ObservationSize:
-    return 15
-
   @property
   def mj_model(self) -> mujoco.MjModel:
     return self._mj_model

mujoco_playground/_src/dm_control_suite/humanoid.py

@@ -215,10 +215,6 @@ def xml_path(self) -> str:
   def action_size(self) -> int:
     return self.mjx_model.nu
 
-  @property
-  def observation_size(self) -> mjx_env.ObservationSize:
-    return 67
-
   @property
   def mj_model(self) -> mujoco.MjModel:
     return self._mj_model

mujoco_playground/_src/dm_control_suite/pendulum.py

@@ -132,10 +132,6 @@ def xml_path(self) -> str:
   def action_size(self) -> int:
     return self.mjx_model.nu
 
-  @property
-  def observation_size(self) -> mjx_env.ObservationSize:
-    return 3
-
   @property
   def mj_model(self) -> mujoco.MjModel:
     return self._mj_model

mujoco_playground/_src/dm_control_suite/point_mass.py

@@ -138,10 +138,6 @@ def xml_path(self) -> str:
   def action_size(self) -> int:
     return self.mjx_model.nu
 
-  @property
-  def observation_size(self) -> mjx_env.ObservationSize:
-    return 4
-
   @property
   def mj_model(self) -> mujoco.MjModel:
     return self._mj_model

mujoco_playground/_src/dm_control_suite/reacher.py

@@ -158,10 +158,6 @@ def xml_path(self) -> str:
   def action_size(self) -> int:
     return self.mjx_model.nu
 
-  @property
-  def observation_size(self) -> mjx_env.ObservationSize:
-    return 6
-
   @property
   def mj_model(self) -> mujoco.MjModel:
     return self._mj_model

mujoco_playground/_src/dm_control_suite/swimmer.py

@@ -233,10 +233,6 @@ def xml_path(self) -> str:
   def action_size(self) -> int:
     return self.mjx_model.nu
 
-  @property
-  def observation_size(self) -> mjx_env.ObservationSize:
-    return 25
-
   @property
   def mj_model(self) -> mujoco.MjModel:
     return self._mj_model

mujoco_playground/_src/dm_control_suite/walker.py

@@ -185,10 +185,6 @@ def xml_path(self) -> str:
   def action_size(self) -> int:
     return self.mjx_model.nu
 
-  @property
-  def observation_size(self) -> mjx_env.ObservationSize:
-    return 94
-
   @property
   def mj_model(self) -> mujoco.MjModel:
     return self._mj_model

mujoco_playground/_src/locomotion/barkour/joystick.py

@@ -457,10 +457,6 @@ def xml_path(self) -> str:
   def action_size(self) -> int:
     return self.mjx_model.nu
 
-  @property
-  def observation_size(self) -> mjx_env.ObservationSize:
-    return 465
-
   @property
   def mj_model(self) -> mujoco.MjModel:
     return self._mj_model

mujoco_playground/_src/locomotion/berkeley_humanoid/joystick.py

@@ -667,10 +667,3 @@ def sample_command(self, rng: jax.Array) -> jax.Array:
         jp.zeros(3),
         jp.hstack([lin_vel_x, lin_vel_y, ang_vel_yaw]),
     )
-
-  @property
-  def observation_size(self) -> mjx_env.ObservationSize:
-    return {
-        "state": (52,),
-        "privileged_state": (114,),
-    }

mujoco_playground/_src/locomotion/g1/joystick.py

@@ -790,10 +790,3 @@ def sample_command(self, rng: jax.Array) -> jax.Array:
         jp.zeros(3),
         jp.hstack([lin_vel_x, lin_vel_y, ang_vel_yaw]),
     )
-
-  @property
-  def observation_size(self) -> mjx_env.ObservationSize:
-    return {
-        "state": (103,),
-        "privileged_state": (216,),
-    }

mujoco_playground/_src/locomotion/go1/getup.py

@@ -367,10 +367,3 @@ def _cost_dof_vel(self, qvel: jax.Array) -> jax.Array:
 
   def _cost_dof_acc(self, qacc: jax.Array) -> jax.Array:
     return jp.sum(jp.square(qacc))
-
-  @property
-  def observation_size(self) -> mjx_env.ObservationSize:
-    return {
-        "state": (42,),
-        "privileged_state": (91,),
-    }

mujoco_playground/_src/locomotion/go1/handstand.py

@@ -374,13 +374,6 @@ def _cost_joint_pos_limits(self, qpos: jax.Array) -> jax.Array:
   def _cost_dof_acc(self, qacc: jax.Array) -> jax.Array:
     return jp.sum(jp.square(qacc))
 
-  @property
-  def observation_size(self) -> mjx_env.ObservationSize:
-    return {
-        "state": (45,),
-        "privileged_state": (94,),
-    }
-
 
 class Footstand(Handstand):
   """Footstand task for Go1."""

mujoco_playground/_src/locomotion/go1/joystick.py

@@ -599,10 +599,3 @@ def sample_command(self, rng: jax.Array, x_k: jax.Array) -> jax.Array:
     w_k = jax.random.bernoulli(w_rng, 0.5, shape=(3,))
     x_kp1 = x_k - w_k * (x_k - y_k * z_k)
     return x_kp1
-
-  @property
-  def observation_size(self) -> mjx_env.ObservationSize:
-    return {
-        "state": (48,),
-        "privileged_state": (123,),
-    }

mujoco_playground/_src/locomotion/h1/inplace_gait_tracking.py

@@ -375,7 +375,3 @@ def _cost_ang_vel(self, global_angvel: jax.Array) -> jax.Array:
   def _cost_lin_vel(self, global_linvel: jax.Array) -> jax.Array:
     # Penalize xy linear velocity.
     return jp.sum(jp.square(global_linvel[:2]))
-
-  @property
-  def observation_size(self) -> mjx_env.ObservationSize:
-    return 186

mujoco_playground/_src/locomotion/h1/joystick_gait_tracking.py

@@ -477,7 +477,3 @@ def sample_command(self, rng: jax.Array) -> jax.Array:
     )
     cmd = jp.hstack([lin_vel_x, lin_vel_y, ang_vel_yaw])
     return cmd
-
-  @property
-  def observation_size(self) -> mjx_env.ObservationSize:
-    return 113

mujoco_playground/_src/locomotion/locomotion_test.py

@@ -36,7 +36,6 @@ def test_can_create_all_environments(self, env_name: str) -> None:
     state = jax.jit(env.step)(state, jp.zeros(env.action_size))
     self.assertIsNotNone(state)
     obs_shape = jax.tree_util.tree_map(lambda x: x.shape, state.obs)
-    obs_shape = obs_shape[0] if isinstance(obs_shape, tuple) else obs_shape
     self.assertEqual(obs_shape, env.observation_size)
     self.assertFalse(jp.isnan(state.data.qpos).any())
 

mujoco_playground/_src/locomotion/op3/joystick.py

@@ -466,7 +466,3 @@ def sample_command(self, rng: jax.Array) -> jax.Array:
 
     cmd = jp.hstack([lin_vel_x, lin_vel_y, ang_vel_yaw])
     return jp.where(jax.random.bernoulli(rng4, 0.1), jp.zeros(3), cmd)
-
-  @property
-  def observation_size(self) -> mjx_env.ObservationSize:
-    return 147

mujoco_playground/_src/locomotion/spot/getup.py

@@ -277,7 +277,3 @@ def _cost_action_rate(
     c1 = jp.sum(jp.square(act - info["last_act"]))
     c2 = jp.sum(jp.square(act - 2 * info["last_act"] + info["last_last_act"]))
     return c1 + c2
-
-  @property
-  def observation_size(self) -> mjx_env.ObservationSize:
-    return 30

mujoco_playground/_src/locomotion/spot/joystick.py

@@ -612,7 +612,3 @@ def sample_command(self, rng: jax.Array) -> jax.Array:
     )
     cmd = jp.hstack([lin_vel_x, lin_vel_y, ang_vel_yaw])
     return jp.where(jax.random.bernoulli(rng4, 0.1), jp.zeros(3), cmd)
-
-  @property
-  def observation_size(self) -> mjx_env.ObservationSize:
-    return {"privileged_state": (167,), "state": (81,)}

mujoco_playground/_src/locomotion/spot/joystick_gait_tracking.py

@@ -386,7 +386,3 @@ def sample_command(self, rng: jax.Array) -> jax.Array:
     )
     cmd = jp.hstack([lin_vel_x, lin_vel_y, ang_vel_yaw])
     return jp.where(jax.random.bernoulli(rng4, 0.1), jp.zeros(3), cmd)
-
-  @property
-  def observation_size(self) -> mjx_env.ObservationSize:
-    return 69

mujoco_playground/_src/manipulation/aloha/handover.py

@@ -283,7 +283,3 @@ def _get_obs(self, data: mjx.Data, info: Dict[str, Any]) -> jax.Array:
     ])
 
     return obs
-
-  @property
-  def observation_size(self) -> int:
-    return 83

mujoco_playground/_src/manipulation/aloha/single_peg_insertion.py

@@ -252,7 +252,3 @@ def _get_reward(
         "peg_z_up": peg_orientation * peg_lift,
         "peg_insertion_reward": peg_insertion_reward,
     }
-
-  @property
-  def observation_size(self) -> mjx_env.ObservationSize:
-    return 82

mujoco_playground/_src/manipulation/franka_emika_panda/open_cabinet.py

@@ -213,7 +213,3 @@ def _get_obs(self, data: mjx.Data, info: dict[str, Any]) -> jax.Array:
     ])
 
     return obs
-
-  @property
-  def observation_size(self) -> mjx_env.ObservationSize:
-    return 55

mujoco_playground/_src/manipulation/franka_emika_panda/pick.py

@@ -230,10 +230,6 @@ def _get_obs(self, data: mjx.Data, info: dict[str, Any]) -> jax.Array:
 
     return obs
 
-  @property
-  def observation_size(self) -> mjx_env.ObservationSize:
-    return 66
-
 
 class PandaPickCubeOrientation(PandaPickCube):
   """Bring a box to a target and orientation."""

mujoco_playground/_src/manipulation/franka_emika_panda/pick_cartesian.py

@@ -433,11 +433,6 @@ def _move_tip(
 
     return new_ctrl, new_tip_pos, no_soln
 
-  @property
-  def observation_size(self) -> mjx_env.ObservationSize:
-    ret = {'pixels/view_0': (64, 64, 3)} if self._vision else 70
-    return ret
-
   @property
   def action_size(self) -> int:
     return 3

mujoco_playground/_src/manipulation/franka_emika_panda_robotiq/push_cube.py

@@ -567,7 +567,3 @@ def _get_single_obs(self, data: mjx.Data, info: dict[str, Any]) -> jax.Array:
   @property
   def action_size(self):
     return 7
-
-  @property
-  def observation_size(self):
-    return 48

mujoco_playground/_src/manipulation/leap_hand/reorient.py

@@ -474,10 +474,6 @@ def get_xfrc(
     data = state.data.replace(xfrc_applied=xfrc)
     return state.replace(data=data)
 
-  @property
-  def observation_size(self) -> mjx_env.ObservationSize:
-    return {"privileged_state": (128,), "state": (57,)}
-
 
 def domain_randomize(model: mjx.Model, rng: jax.Array):
   mj_model = CubeReorient().mj_model

mujoco_playground/_src/manipulation/leap_hand/rotate_z.py

@@ -260,10 +260,6 @@ def _cost_action_rate(
   def _cost_pose(self, joint_angles: jax.Array) -> jax.Array:
     return jp.sum(jp.square(joint_angles - self._default_pose))
 
-  @property
-  def observation_size(self) -> mjx_env.ObservationSize:
-    return {"privileged_state": (105,), "state": (32,)}
-
 
 def domain_randomize(model: mjx.Model, rng: jax.Array):
   mj_model = CubeRotateZAxis().mj_model

mujoco_playground/_src/manipulation/manipulation_test.py

@@ -36,7 +36,6 @@ def test_can_create_all_environments(self, env_name: str) -> None:
     state = jax.jit(env.step)(state, jp.zeros(env.action_size))
     self.assertIsNotNone(state)
     obs_shape = jax.tree_util.tree_map(lambda x: x.shape, state.obs)
-    obs_shape = obs_shape[0] if isinstance(obs_shape, tuple) else obs_shape
     self.assertEqual(obs_shape, env.observation_size)
     self.assertFalse(jp.isnan(state.data.qpos).any())
 

mujoco_playground/_src/mjx_env.py

@@ -270,12 +270,8 @@ def n_substeps(self) -> int:
 
   @property
   def observation_size(self) -> ObservationSize:
-    rng = jax.random.PRNGKey(0)
-    reset_state = self.unwrapped.reset(rng)
-    obs = reset_state.obs
-    if isinstance(obs, jax.Array):
-      return obs.shape[-1]
-    return jax.tree_util.tree_map(lambda x: x.shape, obs)
+    out = jax.eval_shape(self.reset, jax.random.PRNGKey(0))
+    return jax.tree_util.tree_map(lambda x: x.shape, out.obs)
 
   def render(
       self,