diff --git a/configs/main_config.yaml b/configs/main_config.yaml index 745755a..b7cb620 100644 --- a/configs/main_config.yaml +++ b/configs/main_config.yaml @@ -11,6 +11,7 @@ defaults: - morphology: 5_arms_full - arena: default - environment: directed_locomotion + - obs_bounds: default - simulation: default - _self_ diff --git a/configs/obs_bounds/default.yaml b/configs/obs_bounds/default.yaml new file mode 100644 index 0000000..cd70b0b --- /dev/null +++ b/configs/obs_bounds/default.yaml @@ -0,0 +1 @@ +# Defaults provided by dataclass diff --git a/src/brittle_star_project/configs/main_config.py b/src/brittle_star_project/configs/main_config.py index 5a937a7..10fd22e 100644 --- a/src/brittle_star_project/configs/main_config.py +++ b/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) diff --git a/src/brittle_star_project/configs/register_configs.py b/src/brittle_star_project/configs/register_configs.py index b677b31..147da10 100644 --- a/src/brittle_star_project/configs/register_configs.py +++ b/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) diff --git a/src/brittle_star_project/environment/env_config.py b/src/brittle_star_project/environment/env_config.py index 7cb4c21..0af387e 100644 --- a/src/brittle_star_project/environment/env_config.py +++ b/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), + } diff --git a/src/brittle_star_project/environment/obs_processing.py b/src/brittle_star_project/environment/obs_processing.py new file mode 100644 index 0000000..0ef6e3d --- /dev/null +++ b/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))