chore: observation pipeline in simulate

scripts/simulate.py

@@ -27,23 +27,9 @@ from omegaconf import DictConfig, OmegaConf, open_dict
 from brittle_star_project import Backend, BrittleStarEnv, BrittleStarEnvFactory
 from brittle_star_project.configs.main_config import BrittleStarConfig
 from brittle_star_project.configs.register_configs import register_configs
-from brittle_star_project.environment.padded_obs_wrapper import (
+from brittle_star_project.environment.padded_obs_wrapper import compute_padding_masks
-    compute_padding_masks,
+from brittle_star_project.environment.obs_processing import create_obs_processor
-    pad_observation,
+from brittle_star_project.environment.env_config import ObservationBoundsConfig
-)
-
-_ALLOWED_OBS_KEYS = {
-    "joint_position",
-    "joint_velocity",
-    "joint_actuator_force",
-    "actuator_force",
-    "disk_position",
-    "disk_rotation",
-    "disk_linear_velocity",
-    "disk_angular_velocity",
-    "unit_xy_direction_to_target",
-    "xy_distance_to_target",
-}
 
 
 def _dense_layer_sizes_from_params(params: Any) -> list[int]:
@@ -113,28 +99,6 @@ def _maybe_clip_action(
     return np.clip(action, low, high)
 
 
-def _transform_obs_dict(obs_dict: dict[str, Any]) -> jnp.ndarray:
-    """Flatten the env's observation dict into a 1D vector.
-
-    Matches training behavior:
-    - only includes keys in _ALLOWED_OBS_KEYS
-    - iterates keys in sorted order for stable layout
-    - skips empty arrays
-    """
-    parts: list[jnp.ndarray] = []
-    for key in sorted(obs_dict.keys()):
-        if key not in _ALLOWED_OBS_KEYS:
-            continue
-        arr = jnp.asarray(obs_dict[key])
-        if arr.size == 0:
-            continue
-        parts.append(arr.reshape((-1,)))
-
-    if not parts:
-        return jnp.zeros((0,), dtype=jnp.float32)
-    return jnp.concatenate(parts, axis=0)
-
-
 # A minimal policy class to load a CleanRL/Flax checkpoint and run inference.
 class CleanRLPPOPolicy:
     def __init__(
@@ -143,6 +107,7 @@ class CleanRLPPOPolicy:
         sensor_params: Any,
         actor_params: Any,
         action_dim: int,
+        obs_processor: Any,
     ) -> None:
         from brittle_star_project.MLPs.mlps import Actor, GenericDenseLayersWithActivation
 
@@ -155,12 +120,14 @@ class CleanRLPPOPolicy:
             "sensor_params": sensor_params,
             "actor_params": actor_params,
         }
+        self._obs_processor = obs_processor
 
     @staticmethod
     def load(
         path: Path,
         *,
         action_dim: int,
+        obs_processor: Any,
     ) -> "CleanRLPPOPolicy":
         def _get_index(container: Any, idx: int) -> Any:
             if isinstance(container, (list, tuple)):
@@ -280,10 +247,12 @@ class CleanRLPPOPolicy:
             sensor_params=sensor_params,
             actor_params=actor_params,
             action_dim=action_dim,
+            obs_processor=obs_processor,
         )
 
     def act(self, *, observations: dict[str, Any]) -> np.ndarray:
-        obs = _transform_obs_dict(observations)
+        batched_obs = jax.tree.map(lambda x: jnp.asarray(x)[None, ...], observations)
+        obs = self._obs_processor(batched_obs)[0]
         hidden = self._sensor_apply(self._params["sensor_params"], obs)
         mean, _log_std = self._actor_apply(self._params["actor_params"], hidden)
 
@@ -312,7 +281,6 @@ def _rollout_one_episode_headless(
     max_steps: int,
     action_low: np.ndarray | None,
     action_high: np.ndarray | None,
-    padding_masks: dict[str, Any] | None,
 ) -> tuple[float, int, bool, float | None]:
     """Run one rollout up to max_steps.
 
@@ -332,8 +300,6 @@ def _rollout_one_episode_headless(
     steps = 0
     for _ in range(int(max_steps)):
         obs_dict = observations or {}
-        if padding_masks is not None:
-            obs_dict = pad_observation(obs_dict, padding_masks)
 
         action = policy.act(observations=obs_dict)
         action = _maybe_clip_action(action, action_low, action_high)
@@ -369,7 +335,6 @@ def _run_one_episode_viewer(
     max_steps: int | None,
     action_low: np.ndarray | None,
     action_high: np.ndarray | None,
-    padding_masks: dict[str, Any] | None,
 ) -> None:
     import mujoco.viewer
 
@@ -392,8 +357,6 @@ def _run_one_episode_viewer(
             step_start = time.time()
 
             obs_dict = observations or {}
-            if padding_masks is not None:
-                obs_dict = pad_observation(obs_dict, padding_masks)
 
             action = policy.act(observations=obs_dict)
             action = _maybe_clip_action(action, action_low, action_high)
@@ -596,23 +559,29 @@ def main(dict_cfg: DictConfig) -> None:
     # Match training's padded observation layout for amputated morphologies.
     padding_masks = compute_padding_masks(config.morphology.segments_per_arm)
 
-    # Match training's action clipping behavior.
+    training_bounds = ObservationBoundsConfig().to_bounds_dict()
-    action_space = getattr(raw_env, "action_space", None)
+    if trained_cfg_path and "obs_bounds" in trained_cfg:
-    action_low = (
+        try:
-        None if action_space is None else np.asarray(action_space.low, dtype=np.float32).ravel()
+            training_bounds = OmegaConf.to_object(trained_cfg.obs_bounds).to_bounds_dict()
-    )
+        except Exception:
-    action_high = (
+            pass
-        None if action_space is None else np.asarray(action_space.high, dtype=np.float32).ravel()
+    else:
+        training_bounds = config.obs_bounds.to_bounds_dict()
+
+    obs_processor = create_obs_processor(
+        bounds_dict=training_bounds,
+        padding_masks=padding_masks,
     )
 
     # ======= MODEL SETUP =======
     nu = int(state0.mj_model.nu)
-    policy = CleanRLPPOPolicy.load(model_path, action_dim=nu)
+    policy = CleanRLPPOPolicy.load(model_path, action_dim=nu, obs_processor=obs_processor)
 
     # Helpful early failure when configs don't match the checkpoint.
     observations0 = _get_observations(state0)
-    obs0_dict = pad_observation(observations0 or {}, padding_masks)
+    obs0_dict = observations0 or {}
-    env_obs_dim = int(_transform_obs_dict(obs0_dict).shape[0])
+    batched_obs0 = jax.tree.map(lambda x: jnp.asarray(x)[None, ...], obs0_dict)
+    env_obs_dim = int(obs_processor(batched_obs0).shape[1])
     ckpt_obs_dim = _infer_checkpoint_obs_dim(policy)
 
     if ckpt_obs_dim is not None and ckpt_obs_dim != env_obs_dim:
@@ -623,6 +592,15 @@ def main(dict_cfg: DictConfig) -> None:
             "that was used during training."
         )
 
+    # Match training's action clipping behavior.
+    action_space = getattr(raw_env, "action_space", None)
+    action_low = (
+        None if action_space is None else np.asarray(action_space.low, dtype=np.float32).ravel()
+    )
+    action_high = (
+        None if action_space is None else np.asarray(action_space.high, dtype=np.float32).ravel()
+    )
+
     # ======= SIMULATION =======
     headless = bool(config.simulation.headless)
     max_steps = config.simulation.max_steps
@@ -641,7 +619,6 @@ def main(dict_cfg: DictConfig) -> None:
             max_steps=max_steps_i,
             action_low=action_low,
             action_high=action_high,
-            padding_masks=padding_masks,
         )
         final_dist_str = "n/a" if final_dist is None else f"{final_dist:.3f}"
         print(
@@ -670,7 +647,6 @@ def main(dict_cfg: DictConfig) -> None:
             max_steps=max_steps_val,
             action_low=action_low,
             action_high=action_high,
-            padding_masks=padding_masks,
         )
 
     env.close()