feat: extract observations pipeline

configs/main_config.yaml

@@ -11,6 +11,7 @@ defaults:
   - morphology: 5_arms_full
   - arena: default
   - environment: directed_locomotion
+  - obs_bounds: default
   - simulation: default
   - _self_
 

configs/obs_bounds/default.yaml

@@ -0,0 +1 @@
+# Defaults provided by dataclass

src/brittle_star_project/configs/main_config.py

@@ -5,7 +5,12 @@ from brittle_star_project.configs.config_experiment import ExperimentConfig
 from brittle_star_project.configs.config_ppo import PPOConfig
 from brittle_star_project.configs.config_architecture import ArchitectureConfig
 from brittle_star_project.configs.config_simulation import SimulationSettings
-from brittle_star_project.environment.env_config import MorphologyConfig, ArenaConfig, EnvConfig
+from brittle_star_project.environment.env_config import (
+    MorphologyConfig,
+    ArenaConfig,
+    EnvConfig,
+    ObservationBoundsConfig,
+)
 
 
 @dataclass
@@ -25,4 +30,5 @@ class BrittleStarConfig:
     morphology: MorphologyConfig = field(default_factory=MorphologyConfig)
     arena: ArenaConfig = field(default_factory=ArenaConfig)
     environment: EnvConfig = field(default_factory=EnvConfig)
+    obs_bounds: ObservationBoundsConfig = field(default_factory=ObservationBoundsConfig)
     simulation: SimulationSettings = field(default_factory=SimulationSettings)

src/brittle_star_project/configs/register_configs.py

@@ -8,7 +8,12 @@ from brittle_star_project.configs.config_architecture import (
     DecentralizedConfig,
 )
 from brittle_star_project.configs.config_simulation import SimulationSettings
-from brittle_star_project.environment.env_config import MorphologyConfig, ArenaConfig, EnvConfig
+from brittle_star_project.environment.env_config import (
+    MorphologyConfig,
+    ArenaConfig,
+    EnvConfig,
+    ObservationBoundsConfig,
+)
 from brittle_star_project.configs.main_config import BrittleStarConfig
 
 
@@ -37,4 +42,5 @@ def register_configs() -> None:
     cs.store(group="morphology", name="base_morphology", node=MorphologyConfig)
     cs.store(group="arena", name="base_arena", node=ArenaConfig)
     cs.store(group="environment", name="base_environment", node=EnvConfig)
+    cs.store(group="obs_bounds", name="base_obs_bounds", node=ObservationBoundsConfig)
     cs.store(group="simulation", name="base_simulation", node=SimulationSettings)

src/brittle_star_project/environment/env_config.py

@@ -60,3 +60,28 @@ class EnvConfig:
     # Light escape
     # Per docs in upstream env config: integer factors of 200.
     light_perlin_noise_scale: int = 0
+
+
+@dataclass
+class ObservationBoundsConfig:
+    """Physical observation bounds for deterministic min-max normalization."""
+
+    # TODO Inspect empirically observed ranges and update these bounds as needed.
+    joint_position: list[float] = field(default_factory=lambda: [-3.14, 3.14])
+    joint_velocity: list[float] = field(default_factory=lambda: [-20.0, 20.0])
+    joint_actuator_force: list[float] = field(default_factory=lambda: [-5.0, 5.0])
+    segment_contact: list[float] = field(default_factory=lambda: [0.0, 1.0])
+    unit_xy_direction_to_target: list[float] = field(default_factory=lambda: [-1.0, 1.0])
+    xy_distance_to_target: list[float] = field(default_factory=lambda: [0.0, 20.0])
+    disk_z_tilt: list[float] = field(default_factory=lambda: [0.0, 3.141592653589793])
+
+    def to_bounds_dict(self) -> dict[str, tuple[float, float]]:
+        return {
+            "joint_position": tuple(self.joint_position),
+            "joint_velocity": tuple(self.joint_velocity),
+            "joint_actuator_force": tuple(self.joint_actuator_force),
+            "segment_contact": tuple(self.segment_contact),
+            "unit_xy_direction_to_target": tuple(self.unit_xy_direction_to_target),
+            "xy_distance_to_target": tuple(self.xy_distance_to_target),
+            "disk_z_tilt": tuple(self.disk_z_tilt),
+        }

src/brittle_star_project/environment/obs_processing.py

@@ -0,0 +1,83 @@
+import jax
+import jax.numpy as jnp
+from typing import Dict, Tuple, Optional
+
+_JOINT_SCALED_KEYS = frozenset(
+    {
+        "joint_position",
+        "joint_velocity",
+        "joint_actuator_force",
+        "actuator_force",
+    }
+)
+
+_SEGMENT_SCALED_KEYS = frozenset(
+    {
+        "segment_contact",
+    }
+)
+
+
+def create_obs_processor(
+    bounds_dict: Dict[str, Tuple[float, float]], padding_masks: Optional[Dict] = None
+):
+    def _add_derived_features(obs: dict) -> dict:
+        new_obs = dict(obs)
+        if "disk_rotation" in new_obs:
+            rot = new_obs["disk_rotation"]
+            new_obs["disk_z_tilt"] = jnp.sqrt(jnp.pow(rot[0], 2) + jnp.pow(rot[1], 2))
+        return new_obs
+
+    def _normalize_features(obs: dict) -> dict:
+        normalized = {}
+        for key, arr in obs.items():
+            if key in bounds_dict:
+                low, high = bounds_dict[key]
+                if low == -1.0 and high == 1.0:
+                    normalized[key] = jnp.clip(arr, -1.0, 1.0)
+                else:
+                    arr_clipped = jnp.clip(arr, low, high)
+                    normalized[key] = 2.0 * (arr_clipped - low) / (high - low) - 1.0
+            else:
+                normalized[key] = arr
+        return normalized
+
+    def _pad_features(obs: dict) -> dict:
+        padded = {}
+        for key, arr in obs.items():
+            if key in _JOINT_SCALED_KEYS:
+                padded_arr = jnp.zeros(padding_masks["target_size_2x"], dtype=arr.dtype)
+                padded[key] = padded_arr.at[padding_masks["mask_2x"]].set(arr)
+            elif key in _SEGMENT_SCALED_KEYS:
+                padded_arr = jnp.zeros(padding_masks["target_size_1x"], dtype=arr.dtype)
+                padded[key] = padded_arr.at[padding_masks["mask_1x"]].set(arr)
+            else:
+                padded[key] = arr
+        return padded
+
+    def _flatten_features(obs: dict) -> jnp.ndarray:
+        ordered_keys = [
+            "disk_z_tilt",
+            "joint_actuator_force",
+            "joint_position",
+            "joint_velocity",
+            "segment_contact",
+            "unit_xy_direction_to_target",
+            "xy_distance_to_target",
+        ]
+        values = []
+        for key in ordered_keys:
+            if key in obs:
+                arr = jnp.asarray(obs[key]).flatten()
+                if arr.size > 0:
+                    values.append(arr)
+        return jnp.concatenate(values)
+
+    def _process_single(obs_dict: dict) -> jnp.ndarray:
+        processed = _add_derived_features(obs_dict)
+        processed = _normalize_features(processed)
+        if padding_masks is not None:
+            processed = _pad_features(processed)
+        return _flatten_features(processed)
+
+    return jax.jit(jax.vmap(_process_single))