From b29818a144e7591e1728c1f7fa2b1fec0a026647 Mon Sep 17 00:00:00 2001 From: Tibo De Peuter Date: Sat, 25 Apr 2026 17:05:46 +0200 Subject: [PATCH 01/18] docs: Clarify input space in more detail --- docs/design/input_action_spaces.md | 95 +++++++++++++++++++++++++++--- 1 file changed, 88 insertions(+), 7 deletions(-) diff --git a/docs/design/input_action_spaces.md b/docs/design/input_action_spaces.md index 7860e16..8d24d37 100644 --- a/docs/design/input_action_spaces.md +++ b/docs/design/input_action_spaces.md @@ -5,13 +5,24 @@ space (outputs). The control models map these observations directly to physical **Inputs (state space)** -The observation space provides the agent with its current physical state and its objective. +The observation space provides the agent with its current physical state and its navigational objective. With a +decentralized control architecture in mind, we divide these inputs into global and local states. -- Joint positions: the current angles of all joints in the morphology. -- Joint velocities: the current moving speed of the joints. -- Goal vector: instad of just a scalar distance, the goal is represented asa a vector (distance and ange/direction) to +Global inputs, always broadcasted to all nodes: + +- Vertical orientation/tilt: A single, simplified metric representing the tilt or vertical alignment of the agent's + central body/disk, derived from the environment's raw disk rotation 3D vector $[roll, pitch, yaw]$: + $$tilt = sqrt(roll^2 + pitch^2)$$. This represents the deviation from the global Z-axis. +- Goal vector: Instead of just a scalar distance, the goal is represented asa a vector (distance and ange/direction) to the target. +Local inputs, routed directly to specific nodes: + +- Joint positions: The current angles of all joints within the morphology. +- Joint velocities: The current angular velocities of the joints. +- Joint actuator forces: The physical forces currently exerted at each specific joint. +- Segment contact: These values indicate whether each physical segment of the agent is currently touching the ground. + **Outputs (action space)** The action space defines how the agent interacts with the environment. @@ -26,15 +37,34 @@ When designing the state space, we must ask: *Could a human operator perform thi providing only the joint position is insufficient to determine the direction a limb is currently moving. By explicitly including joint velocities, the agent can immediately infer momentum and movement direction without needing to memorize past states. -- Goal Vector (Distance + Angle): Providing only the scalar "distance to the goal" as an input is akin to blindfolding - the robot and asking it to find a target by playing "hot or cold." By providing a full vector, the agent knows - exactly where the target is relative to its current orientation, allowing for directed and efficient locomotion. - Absolute Joint Offsets: The physical Brittle Star robot relies on servo motors (if we were to build this simulated robot), which are inherently position-controlled devices. (Continuous rotation servos exist, but they are less commonly used for joints.) If our network outputted continuous torques (forces), a significant portion of the reinforcement learning process would be wasted on learning low-level PID control dynamics (i.e., how much force to apply to hold a position). Abstracting this away forces the learning algorithm to focus entirely on higher-level gait generation and locomotion. +- Simplified vertical orientation: We drop the full 3D spatial rotation and angular velocity arrays in favor of a + single vertical orientation metric (tilt). For a brittle star moving accross a flat plane, this metric is sufficient + for the agent to sense if it is losing balance or flipping over. +- Force representation: We strictly retain the joint actuator forces and drop the more generic actuator force. Forces + that are explicitly tied to individual joints are significantly easier to route into decentralized, local limb nodes, + which is necessary for our message-passing architecture. +- Goal Vector (Distance + Angle): Providing only the scalar "distance to the goal" as an input is akin to blindfolding + the robot and asking it to find a target by playing "hot or cold." By providing a full vector, the agent knows + exactly where the target is relative to its current orientation, allowing for directed and efficient locomotion. + + The goal representation is explicitly divided into a directional vector and a scalar distance. Keeping the direction + as a normalized unit vector bounds the values to the $[-1, 1]$ range, which stabilizes neural network training. + Providing only a scalar "distance to the goal" would force the agent to learning localized searching behaviors (e.g. + random walks or spiraling) to deduce the direction, drastically increasing the difficulty of the learning task. +- Contact sensors: Segment contact detects external ground interaction and is biologically vital for timing gait + transitions. + + +Specifically, we do not include some available inputs: + +- Global position: Absolute spatial coordinates can cause the agent to overfit to a specific coordinate frame or map, + rather than learning general, adaptable locomotion strategies. ## Limitations and alternatives @@ -50,3 +80,54 @@ Alternative state and action formulations include: - Scalar Goal Distance: Giving the agent only the scalar distance to the target would force it to learn a localized searching behavior (e.g., spiraling or random walks) to determine the correct direction. While biologically plausible for simpler organisms following chemical gradients, it drastically increases the difficulty of the learning task. + +## MuJoCo + +This is what the filtered input vectors look like in MuJoCo, with $J$ joints and $S$ segments: + +- `joint_position`: shape=(J,), dtype=float64 +- `joint_velocity`: shape=(J,), dtype=float64 +- `joint_actuator_force`: shape=(J,), dtype=float64 +- `segment_contact`: shape=(S,), dtype=float64 +- `unit_xy_direction_to_target`: shape=(2,), dtype=float64 +- `xy_distance_to_target`: shape=(1,), dtype=float64 +- `disk_z_tilt`: shape=(1,), dtype=float64, derived from `disk_rotation` + +This brings the entire input space down to $3J + S + 4$ float64's, compared to $4J + S + 15$ float64's for the +unfiltered inputs. + +For reference, these are all the inputs that are available in the MuJoCo environment: + +``` +obs keys: ['joint_position', 'joint_velocity', 'joint_actuator_force', 'actuator_force', 'disk_position', 'disk_rotation', 'disk_linear_velocity', 'disk_angular_velocity', 'tendon_position', 'tendon_velocity', 'segment_contact', 'unit_xy_direction_to_target', 'xy_distance_to_target'] + +raw observations dict: +{'joint_position': array([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.]), + 'joint_velocity': array([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.]), + 'joint_actuator_force': array([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.]), + 'actuator_force': array([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.]), + 'disk_position': array([0. , 0. , 0.11]), + 'disk_rotation': (0.0, -0.0, 0.0), + 'disk_linear_velocity': array([0., 0., 0.]), + 'disk_angular_velocity': array([0., 0., 0.]), + 'tendon_position': array([], dtype=float64), + 'tendon_velocity': array([], dtype=float64), + 'segment_contact': array([0., 0., 0., 0., 0., 0.]), + 'unit_xy_direction_to_target': array([-0.95333378, -0.30191837]), + 'xy_distance_to_target': array([3.])} + +(shapes) +joint_position: shape=(12,), dtype=float64, size=12 +joint_velocity: shape=(12,), dtype=float64, size=12 +joint_actuator_force: shape=(12,), dtype=float64, size=12 +actuator_force: shape=(12,), dtype=float64, size=12 +disk_position: shape=(3,), dtype=float64, size=3 +disk_rotation: shape=(3,), dtype=float64, size=3 +disk_linear_velocity: shape=(3,), dtype=float64, size=3 +disk_angular_velocity: shape=(3,), dtype=float64, size=3 +tendon_position: shape=(0,), dtype=float64, size=0 +tendon_velocity: shape=(0,), dtype=float64, size=0 +segment_contact: shape=(6,), dtype=float64, size=6 +unit_xy_direction_to_target: shape=(2,), dtype=float64, size=2 +xy_distance_to_target: shape=(1,), dtype=float64, size=1 +``` From a1db544e764ad859565467ac6c153e8d24a07e81 Mon Sep 17 00:00:00 2001 From: Tibo De Peuter Date: Tue, 28 Apr 2026 10:43:48 +0200 Subject: [PATCH 02/18] feat: extract observations pipeline --- configs/main_config.yaml | 1 + configs/obs_bounds/default.yaml | 1 + .../configs/main_config.py | 8 +- .../configs/register_configs.py | 8 +- .../environment/env_config.py | 25 ++++++ .../environment/obs_processing.py | 83 +++++++++++++++++++ 6 files changed, 124 insertions(+), 2 deletions(-) create mode 100644 configs/obs_bounds/default.yaml create mode 100644 src/brittle_star_project/environment/obs_processing.py 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)) From 0032496073059af2037a6e3d8ee88d7cf16a037c Mon Sep 17 00:00:00 2001 From: Tibo De Peuter Date: Tue, 28 Apr 2026 10:44:55 +0200 Subject: [PATCH 03/18] docs: describe normalization --- docs/design/input_action_spaces.md | 34 +++++++++++++++++++++++++++--- 1 file changed, 31 insertions(+), 3 deletions(-) diff --git a/docs/design/input_action_spaces.md b/docs/design/input_action_spaces.md index 8d24d37..143dce5 100644 --- a/docs/design/input_action_spaces.md +++ b/docs/design/input_action_spaces.md @@ -10,9 +10,12 @@ decentralized control architecture in mind, we divide these inputs into global a Global inputs, always broadcasted to all nodes: -- Vertical orientation/tilt: A single, simplified metric representing the tilt or vertical alignment of the agent's - central body/disk, derived from the environment's raw disk rotation 3D vector $[roll, pitch, yaw]$: - $$tilt = sqrt(roll^2 + pitch^2)$$. This represents the deviation from the global Z-axis. +- Vertical orientation/tilt: A single, simplified metric representing the tilt/vertical alignment of the agent's + central body/disk, a.k.a. the deviation from the global Z-axis. Its value is derived from the environment's raw disk + rotation 3D vector $[roll, pitch, yaw]$: + $$ + tilt = sqrt(roll^2 + pitch^2) + $$ - Goal vector: Instead of just a scalar distance, the goal is represented asa a vector (distance and ange/direction) to the target. @@ -29,6 +32,16 @@ The action space defines how the agent interacts with the environment. - Joint offsets: *absolute* target positions (offsets) for the joints, i.e. the exact angle the joint should move to. +## Normalization and Scaling + +Both the input (observation) and output (action) spaces are rescaled to the range **$[-1, 1]$**. + +For the input space, all raw physical values (angles, velocities, forces, distances) are normalized based on their +defined physical bounds. If a value exceeds these bounds during simulation, it is clipped to the $[-1, 1]$ range. + +For the output space, the neural network's tanh-activated outputs (which naturally fall in $[-1, 1]$) are linearly +mapped to the physical joint limits defined in the robot's morphology. + ## Rationale When designing the state space, we must ask: *Could a human operator perform this task given only these inputs?* @@ -59,6 +72,17 @@ When designing the state space, we must ask: *Could a human operator perform thi random walks or spiraling) to deduce the direction, drastically increasing the difficulty of the learning task. - Contact sensors: Segment contact detects external ground interaction and is biologically vital for timing gait transitions. +- **Zero-Centered Rescaling ($[-1, 1]$):** Using a zero-centered range is standard best practice for continuous control + tasks. It provides several mathematical and physical advantages: + - **Improved Gradient Flow:** Neural networks optimize faster when inputs are zero-centered. If all inputs were + positive (e.g., $[0, 1]$), the gradients during backpropagation would be forced to the same sign, causing + inefficient "zig-zag" weight updates. + - **Meaningful "Neutral" State:** In robotics, $0.0$ naturally represents a resting state (zero velocity, centered + position, no force). In a $[-1, 1]$ system, this physical rest maps to a neutral $0.0$ signal in the network. + - **Robustness to Amputation:** In this project, amputated limbs are padded with $0.0$. In a $[-1, 1]$ system, this + correctly communicates a "neutral/dead" signal. In a $[0, 1]$ system, $0.0$ would represent the absolute minimum + physical limit, causing the network to misinterpret missing limbs as being at their extreme limits. + Specifically, we do not include some available inputs: @@ -80,6 +104,10 @@ Alternative state and action formulations include: - Scalar Goal Distance: Giving the agent only the scalar distance to the target would force it to learn a localized searching behavior (e.g., spiraling or random walks) to determine the correct direction. While biologically plausible for simpler organisms following chemical gradients, it drastically increases the difficulty of the learning task. +- **$[0, 1]$ Rescaling:** While some domains (like computer vision) use $[0, 1]$ scaling, it is generally avoided in + robotics. Scaling to $[0, 1]$ would mean that a resting joint (velocity = 0) maps to an input of $0.5$. This + constant positive bias forces the network to waste capacity learning to ignore or subtract this baseline signal just + to stand still. Furthermore, it breaks the "dead signal" interpretation of zero-padding used for amputations. ## MuJoCo From 036edf8aab4e85b42d0b38c650d430ad94e6b79d Mon Sep 17 00:00:00 2001 From: Tibo De Peuter Date: Tue, 28 Apr 2026 11:15:16 +0200 Subject: [PATCH 04/18] chore: observation pipeline in training --- .../environment/BrittleStarJaxEnvWrapper.py | 22 +++--- .../trainers/PPOTrainer.py | 79 ++++--------------- 2 files changed, 27 insertions(+), 74 deletions(-) diff --git a/src/brittle_star_project/environment/BrittleStarJaxEnvWrapper.py b/src/brittle_star_project/environment/BrittleStarJaxEnvWrapper.py index a5175a7..b514c11 100644 --- a/src/brittle_star_project/environment/BrittleStarJaxEnvWrapper.py +++ b/src/brittle_star_project/environment/BrittleStarJaxEnvWrapper.py @@ -5,7 +5,7 @@ from experiment_logger import get_logger from .env_config import EnvConfig, MorphologyConfig, ArenaConfig from .env_types import Backend from .factory import BrittleStarEnvFactory -from .padded_obs_wrapper import compute_padding_masks, pad_observations_batched +from .padded_obs_wrapper import compute_padding_masks class BrittleStarJaxEnvWrapper: @@ -48,6 +48,15 @@ class BrittleStarJaxEnvWrapper: def raw(self): return self._env + @property + def padding_masks(self) -> dict: + """Pre-computed boolean masks for amputated limb padding. + + Pass to create_obs_processor so the processor handles padding + after normalization in the correct pipeline order. + """ + return self._padding_masks + @property def single_action_space(self): return self._env.action_space @@ -61,10 +70,6 @@ class BrittleStarJaxEnvWrapper: self._action_rng, env_rng = jax.random.split(jax.random.PRNGKey(seed), 2) env_rngs = jnp.array(jax.random.split(env_rng, self._num_envs)) state = self._vectorized_reset(rng=env_rngs) - - state = state.replace( - observations=pad_observations_batched(state.observations, self._padding_masks) - ) return state def sample_actions(self): @@ -75,12 +80,7 @@ class BrittleStarJaxEnvWrapper: return self._vectorized_action_sample(rng=jnp.array(sub_rngs)) def step(self, state, action): - next_state = self._vectorized_step(state=state, action=action) - - next_state = next_state.replace( - observations=pad_observations_batched(next_state.observations, self._padding_masks) - ) - return next_state + return self._vectorized_step(state=state, action=action) def close(self): self._env.close() diff --git a/src/brittle_star_project/trainers/PPOTrainer.py b/src/brittle_star_project/trainers/PPOTrainer.py index 7cdc1ff..a3e4eac 100644 --- a/src/brittle_star_project/trainers/PPOTrainer.py +++ b/src/brittle_star_project/trainers/PPOTrainer.py @@ -16,6 +16,7 @@ from experiment_logger import get_logger from brittle_star_project.configs.main_config import BrittleStarConfig from brittle_star_project.dataclasses import EpisodeStatistics from brittle_star_project.environment.BrittleStarJaxEnvWrapper import BrittleStarJaxEnvWrapper +from brittle_star_project.environment.obs_processing import create_obs_processor from brittle_star_project.MLPs.mlps import ( Actor, AgentParams, @@ -25,19 +26,6 @@ from brittle_star_project.MLPs.mlps import ( ) from brittle_star_project.ppo import PPO -# TODO: move to config -_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", -} # TODO: clip scaled reward? @@ -65,27 +53,6 @@ def _linear_schedule(count, minibatch_count, update_epochs, num_iterations, lear return learning_rate * frac -@jax.jit -def _normalize_obs(obs, mean, var, eps=1e-8): - return jnp.clip((obs - mean) / jnp.sqrt(var + eps), -10.0, 10.0) - - -@jax.jit -def _convert_obs_dict_to_array(obs_dict: dict) -> jnp.ndarray: - """Convert the raw observation dict → flat array, filtering unwanted keys.""" - - def _filter_and_flatten(o: dict) -> jnp.ndarray: - values = [] - for key in sorted(o.keys()): - if key in _ALLOWED_OBS_KEYS: # TODO: NORMALIZATION or .. of observations?? - v = o[key] - if v.size > 0: - values.append(jnp.asarray(v).flatten()) - return jnp.concatenate(values) - - return jax.vmap(_filter_and_flatten)(obs_dict) - - def _get_action_and_value_noise( sensor: GenericDenseLayersWithActivation, feature_extractor: GenericDenseLayersWithActivation, @@ -168,7 +135,7 @@ def _reward_fn(env_state, next_env_state): return jnp.where(next_env_state.terminated, 50.0, clipped_env_reward - penalty) -def _step_env_wrapped(episode_stats, env_state, action, env_step_fn): +def _step_env_wrapped(episode_stats, env_state, action, env_step_fn, obs_processor): next_env_state = env_step_fn(env_state, action) reward = _reward_fn(env_state, next_env_state) @@ -192,7 +159,7 @@ def _step_env_wrapped(episode_stats, env_state, action, env_step_fn): return ( episode_stats, next_env_state, - (_convert_obs_dict_to_array(next_env_state.observations), reward, done), + (obs_processor(next_env_state.observations), reward, done), ) @@ -303,6 +270,12 @@ class PPOTrainer: self.key = jax.random.PRNGKey(self.experiment.seed) + # Build the centralized observation processor: derive -> normalize -> pad -> flatten. + self.obs_processor = create_obs_processor( + bounds_dict=self.cfg.obs_bounds.to_bounds_dict(), + padding_masks=self.env.padding_masks, + ) + self.sensor, self.feature_extractor, self.actor, self.critic = self._init_agent() self.sensor.apply = jax.jit(self.sensor.apply) self.feature_extractor.apply = jax.jit(self.feature_extractor.apply) @@ -316,7 +289,11 @@ class PPOTrainer: partial( _rollout_jit, max_steps=self.ppo.num_steps, - step_env_fn=partial(_step_env_wrapped, env_step_fn=self.env.step), + step_env_fn=partial( + _step_env_wrapped, + env_step_fn=self.env.step, + obs_processor=self.obs_processor, + ), sensor=self.sensor, feature_extractor=self.feature_extractor, actor=self.actor, @@ -367,10 +344,7 @@ class PPOTrainer: ) dummy_reset = self.env.reset(seed=0) - sample_obs = _convert_obs_dict_to_array(dummy_reset.observations)[0] # take first env - self.obs_mean = jnp.zeros((len(sample_obs),)) - self.obs_var = jnp.ones((len(sample_obs),)) - self.obs_count = 1e-4 + sample_obs = self.obs_processor(dummy_reset.observations)[0] # take first env sensor_params = self.sensor.init(sensor_key, sample_obs) feature_extractor_params = self.feature_extractor.init(feature_extractor_key, sample_obs) actor_params = self.actor.init(actor_key, self.sensor.apply(sensor_params, sample_obs)) @@ -410,25 +384,6 @@ class PPOTrainer: returned_episode_lengths=jnp.zeros(self.ppo.num_envs, dtype=jnp.int32), ) - def _update_obs_stats(self, obs: jnp.ndarray): - batch_mean = jnp.mean(obs, axis=0) - batch_var = jnp.var(obs, axis=0) - batch_count = obs.shape[0] - - delta = batch_mean - self.obs_mean - total_count = self.obs_count + batch_count - - new_mean = self.obs_mean + delta * batch_count / total_count - - m_a = self.obs_var * self.obs_count - m_b = batch_var * batch_count - M2 = m_a + m_b + delta**2 * self.obs_count * batch_count / total_count - new_var = M2 / total_count - - self.obs_mean = new_mean - self.obs_var = new_var - self.obs_count = total_count - def _rollout(self, env_state, next_obs, next_done) -> tuple[Any, ...]: return self._rollout_jit( self.agent_state, @@ -594,7 +549,7 @@ class PPOTrainer: self.logger.log_non_interactive(f"Initial reset started: {time.ctime()}") env_state = self.env.reset(seed=self.experiment.seed) - next_obs = _convert_obs_dict_to_array(env_state.observations) + next_obs = self.obs_processor(env_state.observations) next_done = jnp.zeros(self.ppo.num_envs, dtype=jnp.bool_) self.logger.log_non_interactive(f"Initial reset completed: {time.ctime()}") @@ -609,8 +564,6 @@ class PPOTrainer: env_state, next_obs, next_done, training_measurements, storage = self._step( env_state, next_obs, next_done, iteration=iteration ) - self._update_obs_stats(next_obs) - next_obs = _normalize_obs(next_obs, self.obs_mean, self.obs_var) global_step += self.ppo.num_steps * self.ppo.num_envs self._log( From a3a1f3643f84210936cd539ccd619d1dab0f938b Mon Sep 17 00:00:00 2001 From: Tibo De Peuter Date: Tue, 28 Apr 2026 11:40:28 +0200 Subject: [PATCH 05/18] chore: observation pipeline in simulate --- scripts/simulate.py | 92 +++++++++++++++++---------------------------- 1 file changed, 34 insertions(+), 58 deletions(-) diff --git a/scripts/simulate.py b/scripts/simulate.py index a9a51f4..0e6d01b 100644 --- a/scripts/simulate.py +++ b/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 ( - compute_padding_masks, - pad_observation, -) - -_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", -} +from brittle_star_project.environment.padded_obs_wrapper import compute_padding_masks +from brittle_star_project.environment.obs_processing import create_obs_processor +from brittle_star_project.environment.env_config import ObservationBoundsConfig 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. - 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() + training_bounds = ObservationBoundsConfig().to_bounds_dict() + if trained_cfg_path and "obs_bounds" in trained_cfg: + try: + training_bounds = OmegaConf.to_object(trained_cfg.obs_bounds).to_bounds_dict() + except Exception: + pass + 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) - env_obs_dim = int(_transform_obs_dict(obs0_dict).shape[0]) + obs0_dict = observations0 or {} + 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() From 8f2a5d25edd7ef36faddd4496f6b8ed76deb4a6f Mon Sep 17 00:00:00 2001 From: Tibo De Peuter Date: Tue, 28 Apr 2026 13:17:24 +0200 Subject: [PATCH 06/18] feat(simulate): metadata config loading --- configs/simulation/default.yaml | 8 +- scripts/simulate.py | 557 ++++++------------ .../configs/config_simulation.py | 10 +- .../environment/__init__.py | 7 +- .../environment/padded_obs_wrapper.py | 84 +-- 5 files changed, 190 insertions(+), 476 deletions(-) diff --git a/configs/simulation/default.yaml b/configs/simulation/default.yaml index 84694e0..390cd46 100644 --- a/configs/simulation/default.yaml +++ b/configs/simulation/default.yaml @@ -9,7 +9,7 @@ headless: false # In headless mode this is required; in viewer mode null means "infinite". max_steps: null -# Optional: path to the Hydra config.yaml used during training. -# When set, scripts/simulate.py will use it to default morphology/arena/environment/architecture -# to match training (unless you explicitly override those keys via CLI). -trained_config_path: null +# Optional: override morphology for amputation experiments. +# Points to a morphology config YAML file (e.g. configs/morphology/3_arms.yaml). +# If null, the training morphology from the model's metadata is used. +morphology_override: null diff --git a/scripts/simulate.py b/scripts/simulate.py index 0e6d01b..bee6214 100644 --- a/scripts/simulate.py +++ b/scripts/simulate.py @@ -1,11 +1,10 @@ """Simulate a trained policy in the MuJoCo viewer. -Uses Hydra to load the same BrittleStarConfig that was used during training. -Override settings via CLI, e.g.: - python scripts/simulate.py morphology=3_arms - -To replay a run using the *exact* Hydra config used during training, pass: - python scripts/simulate.py simulation.trained_config_path=runs/.../.hydra/config.yaml \ +Automatically extracts the training configuration (morphology, environment, etc.) +from the sidecar metadata YAML file to ensure simulation perfectly matches training. +Override simulation settings via CLI, e.g.: + uv run scripts/simulate.py \ + simulation.morphology_override=config/morphology/3_arms.yaml \ simulation.model_path=runs/.../final_model.flax """ @@ -22,85 +21,24 @@ import jax import jax.numpy as jnp import numpy as np import yaml -from omegaconf import DictConfig, OmegaConf, open_dict +from omegaconf import DictConfig, OmegaConf 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 compute_padding_masks from brittle_star_project.environment.obs_processing import create_obs_processor -from brittle_star_project.environment.env_config import ObservationBoundsConfig +from brittle_star_project.environment.env_config import ( + MorphologyConfig, + ArenaConfig, + EnvConfig, + ObservationBoundsConfig, +) -def _dense_layer_sizes_from_params(params: Any) -> list[int]: - """Infer GenericDenseLayersWithActivation.layer_sizes from a Flax params tree.""" +class PolicyAgent: + """Wraps a trained Flax actor for deterministic inference.""" - try: - dense_params = params["params"] - except Exception as exc: - raise ValueError("Unexpected sensor params structure (missing 'params')") from exc - - layer_sizes: list[int] = [] - idx = 0 - while True: - key = f"Dense_{idx}" - if key not in dense_params: - break - kernel = dense_params[key]["kernel"] - layer_sizes.append(int(np.asarray(kernel).shape[1])) - idx += 1 - - if not layer_sizes: - raise ValueError("Could not infer Dense_* layers from sensor params") - return layer_sizes - - -def _infer_action_dim_from_actor_params(params: Any) -> int | None: - """Best-effort infer action_dim from a Flax Actor params tree.""" - - try: - dense0 = params["params"]["Dense_0"] - bias = dense0.get("bias") - kernel = dense0.get("kernel") - except Exception: - return None - - if bias is not None: - try: - return int(np.asarray(bias).shape[0]) - except Exception: - return None - - if kernel is not None: - try: - return int(np.asarray(kernel).shape[1]) - except Exception: - return None - - return None - - -def _has_cli_override(overrides: list[str], key: str) -> bool: - prefixes = (f"{key}=", f"{key}.", f"+{key}=", f"+{key}.") - return any(str(o).startswith(prefixes) for o in overrides) - - -def _maybe_clip_action( - action: np.ndarray, - low: np.ndarray | None, - high: np.ndarray | None, -) -> np.ndarray: - if low is None or high is None: - return action - low = np.asarray(low, dtype=np.float32).ravel() - high = np.asarray(high, dtype=np.float32).ravel() - if low.shape != action.shape or high.shape != action.shape: - return action - return np.clip(action, low, high) - - -# A minimal policy class to load a CleanRL/Flax checkpoint and run inference. -class CleanRLPPOPolicy: def __init__( self, *, @@ -111,7 +49,28 @@ class CleanRLPPOPolicy: ) -> None: from brittle_star_project.MLPs.mlps import Actor, GenericDenseLayersWithActivation - layer_sizes = _dense_layer_sizes_from_params(sensor_params) + # Infer layer sizes from params + try: + dense_params = ( + sensor_params.get("params", {}) + if isinstance(sensor_params, dict) + else sensor_params["params"] + ) + except Exception: + dense_params = sensor_params + + layer_sizes = [] + idx = 0 + while True: + key = f"Dense_{idx}" + if key not in dense_params: + break + layer_sizes.append(int(np.asarray(dense_params[key]["kernel"]).shape[1])) + idx += 1 + + if not layer_sizes: + raise ValueError("Could not infer Dense_* layers from sensor params") + self._sensor = GenericDenseLayersWithActivation(layer_sizes=layer_sizes) self._actor = Actor(action_dim=action_dim) self._sensor_apply = jax.jit(self._sensor.apply) @@ -128,122 +87,31 @@ class CleanRLPPOPolicy: *, action_dim: int, obs_processor: Any, - ) -> "CleanRLPPOPolicy": - def _get_index(container: Any, idx: int) -> Any: - if isinstance(container, (list, tuple)): - return container[idx] - if isinstance(container, dict): - return container.get(idx, container.get(str(idx))) - raise KeyError(idx) - - def _looks_like_indexed_dict(container: Any) -> bool: - return ( - isinstance(container, dict) - and container - and all(str(k).isdigit() for k in container.keys()) - ) - - def _parse_checkpoint(restored_obj: Any) -> tuple[Any, Any, Any, Any, Any]: - """Extract checkpoint parts. - - Returns (config_dict, sensor_params, actor_params, critic_params, - feature_extractor_params). - - PPOTrainer saves: - flax.serialization.to_bytes([ - config_dict, - [sensor_params, actor_params, critic_params, feature_extractor_params], - ]) - - msgpack_restore() may restore lists as dicts keyed by string indices - ("0", "1", ...), so we accept both shapes. - """ - - cfg_part: Any | None = None - params_part: Any = restored_obj - - if isinstance(restored_obj, (list, tuple)) and len(restored_obj) >= 2: - cfg_part = restored_obj[0] - params_part = restored_obj[1] - elif _looks_like_indexed_dict(restored_obj) and ( - "0" in restored_obj or "1" in restored_obj - ): - cfg_part = restored_obj.get("0", restored_obj.get(0)) - params_part = restored_obj.get("1", restored_obj.get(1)) - - if _looks_like_indexed_dict(params_part): - sensor_params = _get_index(params_part, 0) - actor_params = _get_index(params_part, 1) - critic_params = _get_index(params_part, 2) - feature_extractor_params = _get_index(params_part, 3) - if sensor_params is None or actor_params is None: - raise ValueError("Missing required params in checkpoint") - return ( - cfg_part, - sensor_params, - actor_params, - critic_params, - feature_extractor_params, - ) - - if isinstance(params_part, (list, tuple)) and len(params_part) >= 2: - sensor_params = params_part[0] - actor_params = params_part[1] - critic_params = params_part[2] if len(params_part) >= 3 else None - feature_extractor_params = params_part[3] if len(params_part) >= 4 else None - return ( - cfg_part, - sensor_params, - actor_params, - critic_params, - feature_extractor_params, - ) - - # Accept a plain dict-shaped Flax params mapping commonly produced - # by saving `agent_state.params` directly. Typical keys are - # 'sensor_params' and 'actor_params', or sometimes nested under 'params'. - if isinstance(restored_obj, dict): - # Top-level params dict - params_sub = restored_obj.get("params", {}) - sensor_params = restored_obj.get("sensor_params") or params_sub.get("sensor_params") - actor_params = restored_obj.get("actor_params") or params_sub.get("actor_params") - critic_params = restored_obj.get("critic_params") or params_sub.get("critic_params") - feature_extractor_params = restored_obj.get( - "feature_extractor_params" - ) or params_sub.get("feature_extractor_params") - # Some checkpoints only save actor+sensor as top-level - if sensor_params is not None and actor_params is not None: - return ( - cfg_part, - sensor_params, - actor_params, - critic_params, - feature_extractor_params, - ) - - raise ValueError( - f"Unexpected checkpoint structure in {path}. " - "Expected [config_dict, [sensor_params, actor_params, critic_params, " - "feature_extractor_params]] or an equivalent dict-indexed variant." - ) - + ) -> "PolicyAgent": payload = path.read_bytes() restored = flax.serialization.msgpack_restore(payload) - _cfg_dict, sensor_params, actor_params, _critic_params, _feature_extractor_params = ( - _parse_checkpoint(restored) - ) - ckpt_action_dim = _infer_action_dim_from_actor_params(actor_params) - if ckpt_action_dim is not None and ckpt_action_dim != action_dim: - raise ValueError( - "Checkpoint/env mismatch: " - f"checkpoint expects action_dim={ckpt_action_dim}, " - f"env provides action_dim={action_dim}. " - "Use the same Hydra config (morphology/arena/environment) " - "that was used during training." - ) + sensor_params = None + actor_params = None - return CleanRLPPOPolicy( + # Extract params from restored checkpoint + if isinstance(restored, dict): + params_sub = restored.get("params", {}) + sensor_params = restored.get("sensor_params") or params_sub.get("sensor_params") + actor_params = restored.get("actor_params") or params_sub.get("actor_params") + elif isinstance(restored, (list, tuple)) and len(restored) >= 2: + params_part = restored[1] + if isinstance(params_part, dict): + sensor_params = params_part.get("0", params_part.get(0)) + actor_params = params_part.get("1", params_part.get(1)) + elif isinstance(params_part, (list, tuple)) and len(params_part) >= 2: + sensor_params = params_part[0] + actor_params = params_part[1] + + if sensor_params is None or actor_params is None: + raise ValueError(f"Could not extract sensor and actor params from checkpoint: {path}") + + return PolicyAgent( sensor_params=sensor_params, actor_params=actor_params, action_dim=action_dim, @@ -273,23 +141,30 @@ def _target_reached(*, state: Any) -> bool: return bool(getattr(state, "terminated", False) or getattr(state, "truncated", False)) -def _rollout_one_episode_headless( +def _maybe_clip_action( + action: np.ndarray, + low: np.ndarray | None, + high: np.ndarray | None, +) -> np.ndarray: + if low is None or high is None: + return action + low = np.asarray(low, dtype=np.float32).ravel() + high = np.asarray(high, dtype=np.float32).ravel() + if low.shape != action.shape or high.shape != action.shape: + return action + return np.clip(action, low, high) + + +def _rollout_headless( *, env: BrittleStarEnv, - policy: CleanRLPPOPolicy, + policy: PolicyAgent, seed: int, max_steps: int, action_low: np.ndarray | None, action_high: np.ndarray | None, + action_mask: np.ndarray | None = None, ) -> tuple[float, int, bool, float | None]: - """Run one rollout up to max_steps. - - Returns (return, length, reached_target, final_xy_dist). - - Note: In the MJC backend, the raw env reward can be 0.0; we compute a simple - progress reward based on xy_distance_to_target. - """ - state = env.reset(seed=seed) ep_return = 0.0 @@ -302,12 +177,10 @@ def _rollout_one_episode_headless( obs_dict = observations or {} action = policy.act(observations=obs_dict) + if action_mask is not None: + action = action[action_mask] action = _maybe_clip_action(action, action_low, action_high) - nu = int(state.mj_model.nu) - if nu > 0 and action.shape != (nu,): - raise ValueError(f"Policy returned action shape {action.shape}, expected ({nu},)") - state = env.step(state=state, action=action) steps += 1 @@ -325,23 +198,23 @@ def _rollout_one_episode_headless( return ep_return, steps, reached_target, final_dist -def _run_one_episode_viewer( +def _rollout_viewer( *, env: BrittleStarEnv, - policy: CleanRLPPOPolicy, + policy: PolicyAgent, seed: int, state: Any, control_dt: float, max_steps: int | None, action_low: np.ndarray | None, action_high: np.ndarray | None, + action_mask: np.ndarray | None = None, ) -> None: import mujoco.viewer model = state.mj_model data = state.mj_data - _ = int(seed) episode_return = 0.0 observations = _get_observations(state) prev_dist = _get_xy_distance_to_target(observations) @@ -359,11 +232,9 @@ def _run_one_episode_viewer( obs_dict = observations or {} action = policy.act(observations=obs_dict) + if action_mask is not None: + action = action[action_mask] action = _maybe_clip_action(action, action_low, action_high) - if model.nu > 0 and action.shape != (int(model.nu),): - raise ValueError( - f"Policy returned action shape {action.shape}, expected ({int(model.nu)},)" - ) # The passive viewer runs a GUI thread; protect MuJoCo state mutation. with viewer.lock(): @@ -398,199 +269,117 @@ def _run_one_episode_viewer( ) -def _infer_checkpoint_obs_dim(policy: CleanRLPPOPolicy) -> int | None: - """Best-effort read of the first Dense kernel input dim (obs dim).""" - - try: - kernel = policy._params["sensor_params"]["params"]["Dense_0"]["kernel"] - return int(getattr(kernel, "shape")[0]) - except Exception: - return None - - -def _load_trained_config(path: Path) -> DictConfig: - """Load a trained config YAML. - - Supports both: - - Hydra's run config (e.g. runs/.../.hydra/config.yaml) - - This project's logger metadata YAMLs, which may contain - ``!!python/object/apply:...`` tags for Enums. - - For safety, we *do not* execute Python constructors from YAML; we only - treat these tags as data and extract their scalar arguments. - """ - - if not path.exists(): - raise FileNotFoundError(f"trained_config_path does not exist: '{path}'.") - if not path.is_file(): - raise ValueError(f"trained_config_path must be a file, got: '{path}'.") - - try: - return OmegaConf.load(path) - except Exception as exc: - python_apply_prefix = "tag:yaml.org,2002:python/object/apply:" - - class _SafeLoaderWithPythonApply(yaml.SafeLoader): - pass - - def _construct_python_apply( - loader: yaml.SafeLoader, - _tag_suffix: str, - node: yaml.Node, - ) -> Any: - if isinstance(node, yaml.SequenceNode): - seq = loader.construct_sequence(node) - if len(seq) == 1: - return seq[0] - return seq - if isinstance(node, yaml.MappingNode): - return loader.construct_mapping(node) - return loader.construct_scalar(node) - - _SafeLoaderWithPythonApply.add_multi_constructor( - python_apply_prefix, _construct_python_apply +def _load_metadata_yaml(model_path: Path) -> dict: + """Discover and load the sidecar metadata YAML file.""" + metadata_path = model_path.with_name(model_path.stem + "_metadata.yaml") + if not metadata_path.exists(): + raise FileNotFoundError( + f"Could not find metadata YAML for {model_path.name}. Expected it at {metadata_path}" ) - - try: - data = yaml.load(path.read_text(encoding="utf-8"), Loader=_SafeLoaderWithPythonApply) - except Exception as yaml_exc: - raise ValueError( - "Failed to load trained_config_path as YAML. " - "If this is a Hydra run, pass the run's '.hydra/config.yaml' file. " - f"Got: '{path}'." - ) from yaml_exc - - if not isinstance(data, dict): - raise ValueError( - "trained_config_path must contain a YAML mapping (dict-like) at the root. " - f"Got type={type(data).__name__} from '{path}'." - ) from exc - - # Normalize known enum-like strings to their Enum *names* so OmegaConf's - # structured config merge behaves like the normal Hydra config. - from brittle_star_project.environment.env_types import Task - - env_cfg = data.get("environment") - if isinstance(env_cfg, dict) and isinstance(env_cfg.get("task"), str): - task_str = str(env_cfg["task"]) - try: - env_cfg["task"] = Task[task_str].name - except Exception: - try: - env_cfg["task"] = Task(task_str).name - except Exception: - pass - - return OmegaConf.create(data) + with open(metadata_path, "r") as f: + return yaml.safe_load(f) @hydra.main(config_path="../configs", config_name="main_config", version_base="1.3") def main(dict_cfg: DictConfig) -> None: - # Compose against the structured schema first, so missing keys are validated. - cfg = OmegaConf.merge(OmegaConf.structured(BrittleStarConfig), dict_cfg) + # 1. Hydra composes ONLY SimulationSettings + cfg = OmegaConf.to_object(OmegaConf.merge(OmegaConf.structured(BrittleStarConfig), dict_cfg)) + sim_cfg = cfg.simulation - # Optional: override env-defining sections (morphology/arena/environment/architecture) - # using the exact Hydra config that was used for training. - trained_cfg_path = cfg.simulation.trained_config_path - if trained_cfg_path: - overrides_raw = OmegaConf.select(cfg, "hydra.overrides.task") or [] - overrides = [str(o) for o in overrides_raw] - - trained_cfg_path_abs = Path(hydra.utils.to_absolute_path(trained_cfg_path)) - trained_cfg = _load_trained_config(trained_cfg_path_abs) - if "hydra" in trained_cfg: - with open_dict(trained_cfg): - del trained_cfg["hydra"] - - with open_dict(cfg): - for key in ("morphology", "arena", "environment", "architecture"): - if key in trained_cfg and not _has_cli_override(overrides, key): - base_node = OmegaConf.select(cfg, key) - override_node = OmegaConf.select(trained_cfg, key) - try: - cfg[key] = OmegaConf.merge(base_node, override_node) - except Exception as exc: - raise ValueError( - "Failed to merge trained config into the active Hydra config. " - f"Key={key!r}, trained_config_path='{trained_cfg_path_abs}'." - ) from exc - - # Convert DictConfig to structured dataclass. - config: BrittleStarConfig = OmegaConf.to_object(cfg) - - backend = Backend.MJC - seed = int(config.experiment.seed) - - if getattr(config.architecture, "name", None) != "centralized": - raise ValueError( - "simulate.py currently only supports architecture=centralized. " - f"Got architecture.name={getattr(config.architecture, 'name', None)!r}. " - "(Training supports decentralized, but simulation wiring for it isn't implemented.)" - ) - - model_path_str = config.simulation.model_path + model_path_str = sim_cfg.model_path if model_path_str is None: raise ValueError( "simulation.model_path must be set to a .flax checkpoint (e.g. final_model.flax)" ) - # Hydra chdir changes CWD; resolve relative paths relative to the invocation. model_path = Path(hydra.utils.to_absolute_path(model_path_str)) if model_path.suffix != ".flax": raise ValueError(f"Expected a '.flax' checkpoint, got '{model_path.name}'.") - # ======= ENVIRONMENT SETUP ======= + # 2. Discover + load sidecar metadata YAML + metadata = _load_metadata_yaml(model_path) + + # 3. Reconstruct typed configs from metadata + trained_morphology = OmegaConf.to_object( + OmegaConf.merge(OmegaConf.structured(MorphologyConfig), metadata.get("morphology", {})) + ) + trained_arena = OmegaConf.to_object( + OmegaConf.merge(OmegaConf.structured(ArenaConfig), metadata.get("arena", {})) + ) + + env_dict = metadata.get("environment", {}) + if isinstance(env_dict.get("task"), str): + from brittle_star_project.environment.env_types import Task + + try: + env_dict["task"] = Task[env_dict["task"]].name + except Exception: + try: + env_dict["task"] = Task(env_dict["task"]).name + except Exception: + pass + + trained_environment = OmegaConf.to_object( + OmegaConf.merge(OmegaConf.structured(EnvConfig), env_dict) + ) + trained_obs_bounds = OmegaConf.to_object( + OmegaConf.merge( + OmegaConf.structured(ObservationBoundsConfig), metadata.get("obs_bounds", {}) + ) + ) + + # 4. Determine environment morphology + if sim_cfg.morphology_override is not None: + override_path = Path(hydra.utils.to_absolute_path(sim_cfg.morphology_override)) + if not override_path.exists(): + raise FileNotFoundError(f"Could not find morphology override YAML at {override_path}") + with open(override_path, "r") as f: + override_dict = yaml.safe_load(f) + env_morphology = OmegaConf.to_object( + OmegaConf.merge(OmegaConf.structured(MorphologyConfig), override_dict) + ) + else: + env_morphology = trained_morphology + + # 5. Build obs_processor with TRAINING morphology padding masks always + padding_masks = compute_padding_masks( + segments_per_arm=env_morphology.segments_per_arm, + ) + obs_processor = create_obs_processor( + bounds_dict=trained_obs_bounds.to_bounds_dict(), + padding_masks=padding_masks, + ) + + # 6. Build environment + backend = Backend.MJC + seed = int(cfg.experiment.seed) + factory = BrittleStarEnvFactory() raw_env = factory.create_environment( backend, - config.morphology, - config.arena, - config.environment, + env_morphology, + trained_arena, + trained_environment, ) env = BrittleStarEnv( raw_env, backend=backend, - config=config.environment, - morphology_config=config.morphology, + config=trained_environment, + morphology_config=env_morphology, ) state0 = env.reset(seed=seed) - # Match training's padded observation layout for amputated morphologies. - padding_masks = compute_padding_masks(config.morphology.segments_per_arm) + # Calculate the action dimension the model was trained with + trained_action_dim = sum(trained_morphology.segments_per_arm) * 2 - training_bounds = ObservationBoundsConfig().to_bounds_dict() - if trained_cfg_path and "obs_bounds" in trained_cfg: - try: - training_bounds = OmegaConf.to_object(trained_cfg.obs_bounds).to_bounds_dict() - except Exception: - pass - else: - training_bounds = config.obs_bounds.to_bounds_dict() - - obs_processor = create_obs_processor( - bounds_dict=training_bounds, - padding_masks=padding_masks, + # 7. Load policy + policy = PolicyAgent.load( + model_path, action_dim=trained_action_dim, obs_processor=obs_processor ) - # ======= MODEL SETUP ======= - nu = int(state0.mj_model.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 = observations0 or {} - 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: - raise ValueError( - "Checkpoint/env mismatch: " - f"checkpoint expects obs_dim={ckpt_obs_dim}, env provides obs_dim={env_obs_dim}. " - "Use the same Hydra config (morphology/arena/environment) " - "that was used during training." - ) + # Convert the JAX boolean mask to a numpy array for easy indexing + action_mask = np.asarray(padding_masks["mask_2x"]) # Match training's action clipping behavior. action_space = getattr(raw_env, "action_space", None) @@ -601,24 +390,26 @@ def main(dict_cfg: DictConfig) -> None: 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 + # 8. Run simulation + headless = bool(sim_cfg.headless) + max_steps = sim_cfg.max_steps if headless: if max_steps is None: raise ValueError("simulation.max_steps is required when simulation.headless=true") + max_steps_i = int(max_steps) if max_steps_i <= 0: raise ValueError("simulation.max_steps must be > 0") - ep_return, ep_len, reached_target, final_dist = _rollout_one_episode_headless( + ep_return, ep_len, reached_target, final_dist = _rollout_headless( env=env, policy=policy, seed=seed, max_steps=max_steps_i, action_low=action_low, action_high=action_high, + action_mask=action_mask, ) final_dist_str = "n/a" if final_dist is None else f"{final_dist:.3f}" print( @@ -627,18 +418,17 @@ def main(dict_cfg: DictConfig) -> None: f"target_reached={reached_target}, final_xy_dist={final_dist_str}" ) else: + max_steps_val = None if max_steps is not None: max_steps_i = int(max_steps) if max_steps_i <= 0: raise ValueError("simulation.max_steps must be > 0") - max_steps_val: int | None = max_steps_i - else: - max_steps_val = None + max_steps_val = max_steps_i model_dt = float(state0.mj_model.opt.timestep) - control_dt = model_dt * float(config.environment.num_physics_steps_per_control_step) + control_dt = model_dt * float(trained_environment.num_physics_steps_per_control_step) - _run_one_episode_viewer( + _rollout_viewer( env=env, policy=policy, seed=seed, @@ -647,6 +437,7 @@ def main(dict_cfg: DictConfig) -> None: max_steps=max_steps_val, action_low=action_low, action_high=action_high, + action_mask=action_mask, ) env.close() diff --git a/src/brittle_star_project/configs/config_simulation.py b/src/brittle_star_project/configs/config_simulation.py index 872ca20..a10682d 100644 --- a/src/brittle_star_project/configs/config_simulation.py +++ b/src/brittle_star_project/configs/config_simulation.py @@ -13,7 +13,9 @@ class SimulationSettings: # If None, viewer mode runs until window closed or target reached. max_steps: Optional[int] = None - # Optional: point to a Hydra config.yaml from a training run (e.g. runs/.../.hydra/config.yaml). - # When set, the simulation script can override - # morphology/arena/environment/architecture to match. - trained_config_path: Optional[str] = None + # Override morphology for amputation experiments. + # When set, the environment uses this morphology instead of the trained one. + # Points to a morphology config YAML file (e.g. configs/morphology/3_arms.yaml). + # Observations are padded from the override morphology UP TO the training + # morphology's shape via compute_padding_masks(override, reference=training). + morphology_override: Optional[str] = None diff --git a/src/brittle_star_project/environment/__init__.py b/src/brittle_star_project/environment/__init__.py index 78896cb..77d106c 100644 --- a/src/brittle_star_project/environment/__init__.py +++ b/src/brittle_star_project/environment/__init__.py @@ -1,7 +1,9 @@ from .env_config import ArenaConfig, EnvConfig, MorphologyConfig from .env_types import Backend, Task -from .env_wrapper import BrittleStarEnv, StepResult +from .env_wrapper import BrittleStarEnv from .factory import BrittleStarEnvFactory +from .obs_processing import create_obs_processor +from .padded_obs_wrapper import compute_padding_masks __all__ = [ "ArenaConfig", @@ -10,6 +12,7 @@ __all__ = [ "Backend", "Task", "BrittleStarEnv", - "StepResult", "BrittleStarEnvFactory", + "create_obs_processor", + "compute_padding_masks", ] diff --git a/src/brittle_star_project/environment/padded_obs_wrapper.py b/src/brittle_star_project/environment/padded_obs_wrapper.py index ae64de4..ea00713 100644 --- a/src/brittle_star_project/environment/padded_obs_wrapper.py +++ b/src/brittle_star_project/environment/padded_obs_wrapper.py @@ -1,35 +1,11 @@ -"""Observation padding wrapper for amputated brittle star morphologies. - -When using a centralized controller, the global observation vector must remain -a constant size regardless of how many segments are amputated. This wrapper pads -the observation dictionary values with zeros using spatial insertion so that the -flattened observation maintains the correct physical mapping to the neural network. -""" +"""Observation padding masks for amputated brittle star morphologies.""" from __future__ import annotations from typing import Any, Sequence -import jax import jax.numpy as jnp -# Observation keys whose size scales with the number of joints (2 per segment). -_JOINT_SCALED_KEYS = frozenset( - { - "joint_position", - "joint_velocity", - "joint_actuator_force", - "actuator_force", - } -) - -# Observation keys whose size scales with the number of segments (1 per segment). -_SEGMENT_SCALED_KEYS = frozenset( - { - "segment_contact", - } -) - def compute_padding_masks( segments_per_arm: Sequence[int], @@ -60,7 +36,6 @@ def compute_padding_masks( f"actual segments ({actual}) must be between 0 and reference ({ref})." ) # 1x scaling (e.g., contacts: 1 value per segment) - # 1x scaling (e.g., contacts: 1 value per segment) mask_1x.extend([True] * actual + [False] * (ref - actual)) # 2x scaling (e.g., joints: 2 values per segment) mask_2x.extend([True] * (actual * 2) + [False] * ((ref - actual) * 2)) @@ -71,60 +46,3 @@ def compute_padding_masks( "target_size_1x": sum(reference_segments_per_arm), "target_size_2x": sum(reference_segments_per_arm) * 2, } - - -def pad_observation( - obs: dict[str, Any], - masks: dict[str, Any], -) -> dict[str, Any]: - """Pad an observation dict using spatial insertion.""" - padded = {} - for key, value in obs.items(): - padded_dtype = _padding_dtype(value) - if key in _JOINT_SCALED_KEYS: - out = jnp.zeros(masks["target_size_2x"], dtype=padded_dtype) - padded[key] = out.at[masks["mask_2x"]].set(value) - elif key in _SEGMENT_SCALED_KEYS: - out = jnp.zeros(masks["target_size_1x"], dtype=padded_dtype) - padded[key] = out.at[masks["mask_1x"]].set(value) - else: - padded[key] = value - return padded - - -def pad_observations_batched( - obs: dict[str, Any], - masks: dict[str, Any], -) -> dict[str, Any]: - """Pad a batched observation dict (leading batch dimension) using spatial insertion.""" - padded = {} - for key, value in obs.items(): - batch_size = value.shape[0] - padded_dtype = _padding_dtype(value) - if key in _JOINT_SCALED_KEYS: - out = jnp.zeros((batch_size, masks["target_size_2x"]), dtype=padded_dtype) - padded[key] = out.at[:, masks["mask_2x"]].set(value) - elif key in _SEGMENT_SCALED_KEYS: - out = jnp.zeros((batch_size, masks["target_size_1x"]), dtype=padded_dtype) - padded[key] = out.at[:, masks["mask_1x"]].set(value) - else: - padded[key] = value - return padded - - -def _padding_dtype(value: Any) -> jnp.dtype: - """Choose a JAX-safe dtype for padding arrays. - - When JAX x64 is disabled, allocating float64 zeros emits a warning. We - preserve the original dtype whenever it is supported, and otherwise fall - back to float32 for padding buffers. - """ - - dtype = getattr(value, "dtype", None) - if dtype is None: - dtype = jnp.asarray(value).dtype - else: - dtype = jnp.dtype(dtype) - if dtype == jnp.float64 and not jax.config.read("jax_enable_x64"): - return jnp.float32 - return dtype From 9d1e2c9bfff74c291a4f5fd5b50d567c8f4ada6e Mon Sep 17 00:00:00 2001 From: Tibo De Peuter Date: Tue, 28 Apr 2026 13:39:01 +0200 Subject: [PATCH 07/18] refactor: evaluation subpackage --- scripts/simulate.py | 320 ++---------------- src/brittle_star_project/__init__.py | 16 +- .../evaluation/__init__.py | 17 + .../evaluation/checkpoint.py | 105 ++++++ src/brittle_star_project/evaluation/policy.py | 89 +++++ .../evaluation/rollout.py | 163 +++++++++ src/brittle_star_project/render/__init__.py | 3 - src/brittle_star_project/render/renderer.py | 78 ----- 8 files changed, 409 insertions(+), 382 deletions(-) create mode 100644 src/brittle_star_project/evaluation/__init__.py create mode 100644 src/brittle_star_project/evaluation/checkpoint.py create mode 100644 src/brittle_star_project/evaluation/policy.py create mode 100644 src/brittle_star_project/evaluation/rollout.py delete mode 100644 src/brittle_star_project/render/__init__.py delete mode 100644 src/brittle_star_project/render/renderer.py diff --git a/scripts/simulate.py b/scripts/simulate.py index bee6214..7bb2800 100644 --- a/scripts/simulate.py +++ b/scripts/simulate.py @@ -4,280 +4,29 @@ Automatically extracts the training configuration (morphology, environment, etc. from the sidecar metadata YAML file to ensure simulation perfectly matches training. Override simulation settings via CLI, e.g.: uv run scripts/simulate.py \ - simulation.morphology_override=config/morphology/3_arms.yaml \ + simulation.morphology_override=configs/morphology/3_arms.yaml \ simulation.model_path=runs/.../final_model.flax """ from __future__ import annotations -import itertools -import time from pathlib import Path -from typing import Any -import flax import hydra -import jax -import jax.numpy as jnp import numpy as np -import yaml from omegaconf import DictConfig, OmegaConf +import yaml 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 compute_padding_masks from brittle_star_project.environment.obs_processing import create_obs_processor -from brittle_star_project.environment.env_config import ( - MorphologyConfig, - ArenaConfig, - EnvConfig, - ObservationBoundsConfig, -) +from brittle_star_project.environment.env_config import MorphologyConfig - -class PolicyAgent: - """Wraps a trained Flax actor for deterministic inference.""" - - def __init__( - self, - *, - sensor_params: Any, - actor_params: Any, - action_dim: int, - obs_processor: Any, - ) -> None: - from brittle_star_project.MLPs.mlps import Actor, GenericDenseLayersWithActivation - - # Infer layer sizes from params - try: - dense_params = ( - sensor_params.get("params", {}) - if isinstance(sensor_params, dict) - else sensor_params["params"] - ) - except Exception: - dense_params = sensor_params - - layer_sizes = [] - idx = 0 - while True: - key = f"Dense_{idx}" - if key not in dense_params: - break - layer_sizes.append(int(np.asarray(dense_params[key]["kernel"]).shape[1])) - idx += 1 - - if not layer_sizes: - raise ValueError("Could not infer Dense_* layers from sensor params") - - self._sensor = GenericDenseLayersWithActivation(layer_sizes=layer_sizes) - self._actor = Actor(action_dim=action_dim) - self._sensor_apply = jax.jit(self._sensor.apply) - self._actor_apply = jax.jit(self._actor.apply) - self._params = { - "sensor_params": sensor_params, - "actor_params": actor_params, - } - self._obs_processor = obs_processor - - @staticmethod - def load( - path: Path, - *, - action_dim: int, - obs_processor: Any, - ) -> "PolicyAgent": - payload = path.read_bytes() - restored = flax.serialization.msgpack_restore(payload) - - sensor_params = None - actor_params = None - - # Extract params from restored checkpoint - if isinstance(restored, dict): - params_sub = restored.get("params", {}) - sensor_params = restored.get("sensor_params") or params_sub.get("sensor_params") - actor_params = restored.get("actor_params") or params_sub.get("actor_params") - elif isinstance(restored, (list, tuple)) and len(restored) >= 2: - params_part = restored[1] - if isinstance(params_part, dict): - sensor_params = params_part.get("0", params_part.get(0)) - actor_params = params_part.get("1", params_part.get(1)) - elif isinstance(params_part, (list, tuple)) and len(params_part) >= 2: - sensor_params = params_part[0] - actor_params = params_part[1] - - if sensor_params is None or actor_params is None: - raise ValueError(f"Could not extract sensor and actor params from checkpoint: {path}") - - return PolicyAgent( - 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: - 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) - - # Always evaluate with the actor mean. - # (Sampling adds exploration noise, which is useful for training but not for evaluation.) - return np.asarray(mean, dtype=np.float32).ravel() - - -def _get_observations(state: Any) -> dict[str, Any] | None: - return getattr(state, "observations", None) - - -def _get_xy_distance_to_target(observations: dict[str, Any]) -> float | None: - return float(np.asarray(observations["xy_distance_to_target"]).reshape(-1)[0]) - - -def _target_reached(*, state: Any) -> bool: - return bool(getattr(state, "terminated", False) or getattr(state, "truncated", False)) - - -def _maybe_clip_action( - action: np.ndarray, - low: np.ndarray | None, - high: np.ndarray | None, -) -> np.ndarray: - if low is None or high is None: - return action - low = np.asarray(low, dtype=np.float32).ravel() - high = np.asarray(high, dtype=np.float32).ravel() - if low.shape != action.shape or high.shape != action.shape: - return action - return np.clip(action, low, high) - - -def _rollout_headless( - *, - env: BrittleStarEnv, - policy: PolicyAgent, - seed: int, - max_steps: int, - action_low: np.ndarray | None, - action_high: np.ndarray | None, - action_mask: np.ndarray | None = None, -) -> tuple[float, int, bool, float | None]: - state = env.reset(seed=seed) - - ep_return = 0.0 - observations = _get_observations(state) - prev_dist = _get_xy_distance_to_target(observations) - reached_target = _target_reached(state=state) - - steps = 0 - for _ in range(int(max_steps)): - obs_dict = observations or {} - - action = policy.act(observations=obs_dict) - if action_mask is not None: - action = action[action_mask] - action = _maybe_clip_action(action, action_low, action_high) - - state = env.step(state=state, action=action) - steps += 1 - - observations = _get_observations(state) - cur_dist = _get_xy_distance_to_target(observations) - if prev_dist is not None and cur_dist is not None: - ep_return += prev_dist - cur_dist - prev_dist = cur_dist - - reached_target = _target_reached(state=state) - if reached_target: - break - - final_dist = _get_xy_distance_to_target(observations) - return ep_return, steps, reached_target, final_dist - - -def _rollout_viewer( - *, - env: BrittleStarEnv, - policy: PolicyAgent, - seed: int, - state: Any, - control_dt: float, - max_steps: int | None, - action_low: np.ndarray | None, - action_high: np.ndarray | None, - action_mask: np.ndarray | None = None, -) -> None: - import mujoco.viewer - - model = state.mj_model - data = state.mj_data - - episode_return = 0.0 - observations = _get_observations(state) - prev_dist = _get_xy_distance_to_target(observations) - reached_target = _target_reached(state=state) - - steps = 0 - # Use the viewer as a context manager to avoid GLX teardown races. - with mujoco.viewer.launch_passive(model, data) as viewer: - step_iter = range(int(max_steps)) if max_steps is not None else itertools.count() - for _step_idx in step_iter: - if not viewer.is_running(): - break - step_start = time.time() - - obs_dict = observations or {} - - action = policy.act(observations=obs_dict) - if action_mask is not None: - action = action[action_mask] - action = _maybe_clip_action(action, action_low, action_high) - - # The passive viewer runs a GUI thread; protect MuJoCo state mutation. - with viewer.lock(): - state = env.step(state=state, action=action) - - if not viewer.is_running(): - break - viewer.sync() - - steps += 1 - - observations = _get_observations(state) - cur_dist = _get_xy_distance_to_target(observations) - if prev_dist is not None and cur_dist is not None: - episode_return += prev_dist - cur_dist - prev_dist = cur_dist - - reached_target = _target_reached(state=state) - if reached_target: - break - - remaining = control_dt - (time.time() - step_start) - if remaining > 0: - time.sleep(remaining) - - dist = _get_xy_distance_to_target(observations) - dist_str = "n/a" if dist is None else f"{dist:.3f}" - print( - "episode done: " - f"return={episode_return:.6f}, len={steps}, " - f"target_reached={reached_target}, final_xy_dist={dist_str}" - ) - - -def _load_metadata_yaml(model_path: Path) -> dict: - """Discover and load the sidecar metadata YAML file.""" - metadata_path = model_path.with_name(model_path.stem + "_metadata.yaml") - if not metadata_path.exists(): - raise FileNotFoundError( - f"Could not find metadata YAML for {model_path.name}. Expected it at {metadata_path}" - ) - with open(metadata_path, "r") as f: - return yaml.safe_load(f) +from brittle_star_project.evaluation.checkpoint import load_metadata, metadata_to_configs +from brittle_star_project.evaluation.policy import PolicyAgent +from brittle_star_project.evaluation.rollout import rollout_headless, rollout_viewer @hydra.main(config_path="../configs", config_name="main_config", version_base="1.3") @@ -297,36 +46,10 @@ def main(dict_cfg: DictConfig) -> None: raise ValueError(f"Expected a '.flax' checkpoint, got '{model_path.name}'.") # 2. Discover + load sidecar metadata YAML - metadata = _load_metadata_yaml(model_path) + metadata = load_metadata(model_path) # 3. Reconstruct typed configs from metadata - trained_morphology = OmegaConf.to_object( - OmegaConf.merge(OmegaConf.structured(MorphologyConfig), metadata.get("morphology", {})) - ) - trained_arena = OmegaConf.to_object( - OmegaConf.merge(OmegaConf.structured(ArenaConfig), metadata.get("arena", {})) - ) - - env_dict = metadata.get("environment", {}) - if isinstance(env_dict.get("task"), str): - from brittle_star_project.environment.env_types import Task - - try: - env_dict["task"] = Task[env_dict["task"]].name - except Exception: - try: - env_dict["task"] = Task(env_dict["task"]).name - except Exception: - pass - - trained_environment = OmegaConf.to_object( - OmegaConf.merge(OmegaConf.structured(EnvConfig), env_dict) - ) - trained_obs_bounds = OmegaConf.to_object( - OmegaConf.merge( - OmegaConf.structured(ObservationBoundsConfig), metadata.get("obs_bounds", {}) - ) - ) + training = metadata_to_configs(metadata) # 4. Determine environment morphology if sim_cfg.morphology_override is not None: @@ -339,14 +62,15 @@ def main(dict_cfg: DictConfig) -> None: OmegaConf.merge(OmegaConf.structured(MorphologyConfig), override_dict) ) else: - env_morphology = trained_morphology + env_morphology = training.morphology # 5. Build obs_processor with TRAINING morphology padding masks always padding_masks = compute_padding_masks( segments_per_arm=env_morphology.segments_per_arm, + reference_segments_per_arm=training.morphology.segments_per_arm, ) obs_processor = create_obs_processor( - bounds_dict=trained_obs_bounds.to_bounds_dict(), + bounds_dict=training.obs_bounds.to_bounds_dict(), padding_masks=padding_masks, ) @@ -358,23 +82,23 @@ def main(dict_cfg: DictConfig) -> None: raw_env = factory.create_environment( backend, env_morphology, - trained_arena, - trained_environment, + training.arena, + training.environment, ) env = BrittleStarEnv( raw_env, backend=backend, - config=trained_environment, + config=training.environment, morphology_config=env_morphology, ) state0 = env.reset(seed=seed) # Calculate the action dimension the model was trained with - trained_action_dim = sum(trained_morphology.segments_per_arm) * 2 + trained_action_dim = sum(training.morphology.segments_per_arm) * 2 # 7. Load policy - policy = PolicyAgent.load( + policy = PolicyAgent.from_checkpoint( model_path, action_dim=trained_action_dim, obs_processor=obs_processor ) @@ -402,7 +126,7 @@ def main(dict_cfg: DictConfig) -> None: if max_steps_i <= 0: raise ValueError("simulation.max_steps must be > 0") - ep_return, ep_len, reached_target, final_dist = _rollout_headless( + result = rollout_headless( env=env, policy=policy, seed=seed, @@ -411,11 +135,11 @@ def main(dict_cfg: DictConfig) -> None: action_high=action_high, action_mask=action_mask, ) - final_dist_str = "n/a" if final_dist is None else f"{final_dist:.3f}" + final_dist_str = "n/a" if result.final_xy_dist is None else f"{result.final_xy_dist:.3f}" print( "episode done: " - f"return={ep_return:.6f}, len={ep_len}, " - f"target_reached={reached_target}, final_xy_dist={final_dist_str}" + f"return={result.return_:.6f}, len={result.length}, " + f"target_reached={result.reached_target}, final_xy_dist={final_dist_str}" ) else: max_steps_val = None @@ -426,9 +150,9 @@ def main(dict_cfg: DictConfig) -> None: max_steps_val = max_steps_i model_dt = float(state0.mj_model.opt.timestep) - control_dt = model_dt * float(trained_environment.num_physics_steps_per_control_step) + control_dt = model_dt * float(training.environment.num_physics_steps_per_control_step) - _rollout_viewer( + rollout_viewer( env=env, policy=policy, seed=seed, diff --git a/src/brittle_star_project/__init__.py b/src/brittle_star_project/__init__.py index 3902cbb..4eec766 100644 --- a/src/brittle_star_project/__init__.py +++ b/src/brittle_star_project/__init__.py @@ -2,7 +2,14 @@ from .environment.env_types import Backend, Task from .environment.env_config import ArenaConfig, EnvConfig, MorphologyConfig from .environment.factory import BrittleStarEnvFactory from .environment.env_wrapper import BrittleStarEnv -from .render import simulate_policy, SimulationConfig, ControlPolicy +from .evaluation import ( + PolicyAgent, + ControlPolicy, + load_metadata, + rollout_headless, + rollout_viewer, + EpisodeResult, +) __all__ = [ "ArenaConfig", @@ -12,7 +19,10 @@ __all__ = [ "EnvConfig", "MorphologyConfig", "Task", - "simulate_policy", - "SimulationConfig", + "PolicyAgent", "ControlPolicy", + "load_metadata", + "rollout_headless", + "rollout_viewer", + "EpisodeResult", ] diff --git a/src/brittle_star_project/evaluation/__init__.py b/src/brittle_star_project/evaluation/__init__.py new file mode 100644 index 0000000..73719e6 --- /dev/null +++ b/src/brittle_star_project/evaluation/__init__.py @@ -0,0 +1,17 @@ +from __future__ import annotations + +from .checkpoint import load_metadata, load_params, metadata_to_configs, TrainingConfig +from .policy import PolicyAgent, ControlPolicy +from .rollout import rollout_headless, rollout_viewer, EpisodeResult + +__all__ = [ + "load_metadata", + "load_params", + "metadata_to_configs", + "TrainingConfig", + "PolicyAgent", + "ControlPolicy", + "rollout_headless", + "rollout_viewer", + "EpisodeResult", +] diff --git a/src/brittle_star_project/evaluation/checkpoint.py b/src/brittle_star_project/evaluation/checkpoint.py new file mode 100644 index 0000000..ffbc27a --- /dev/null +++ b/src/brittle_star_project/evaluation/checkpoint.py @@ -0,0 +1,105 @@ +from __future__ import annotations + +import yaml +from dataclasses import dataclass +from pathlib import Path + +import flax +from omegaconf import OmegaConf + +from brittle_star_project.environment.env_config import ( + MorphologyConfig, + ArenaConfig, + EnvConfig, + ObservationBoundsConfig, +) + + +@dataclass +class TrainingConfig: + """Holds typed configurations extracted from a training run's metadata.""" + + morphology: MorphologyConfig + arena: ArenaConfig + environment: EnvConfig + obs_bounds: ObservationBoundsConfig + + +def load_params(path: Path) -> dict: + """Load model parameters from a .flax checkpoint file.""" + payload = path.read_bytes() + restored = flax.serialization.msgpack_restore(payload) + + sensor_params = None + actor_params = None + + # Extract params from restored checkpoint + if isinstance(restored, dict): + params_sub = restored.get("params", {}) + sensor_params = restored.get("sensor_params") or params_sub.get("sensor_params") + actor_params = restored.get("actor_params") or params_sub.get("actor_params") + elif isinstance(restored, (list, tuple)) and len(restored) >= 2: + params_part = restored[1] + if isinstance(params_part, dict): + sensor_params = params_part.get("0", params_part.get(0)) + actor_params = params_part.get("1", params_part.get(1)) + elif isinstance(params_part, (list, tuple)) and len(params_part) >= 2: + sensor_params = params_part[0] + actor_params = params_part[1] + + if sensor_params is None or actor_params is None: + raise ValueError(f"Could not extract sensor and actor params from checkpoint: {path}") + + return { + "sensor_params": sensor_params, + "actor_params": actor_params, + } + + +def load_metadata(model_path: Path) -> dict: + """Discover and load the sidecar metadata YAML file.""" + metadata_path = model_path.with_name(model_path.stem + "_metadata.yaml") + if not metadata_path.exists(): + raise FileNotFoundError( + f"Could not find metadata YAML for {model_path.name}. Expected it at {metadata_path}" + ) + with open(metadata_path, "r") as f: + return yaml.safe_load(f) + + +def metadata_to_configs(metadata: dict) -> TrainingConfig: + """Reconstruct typed configuration objects from a metadata dictionary.""" + trained_morphology = OmegaConf.to_object( + OmegaConf.merge(OmegaConf.structured(MorphologyConfig), metadata.get("morphology", {})) + ) + trained_arena = OmegaConf.to_object( + OmegaConf.merge(OmegaConf.structured(ArenaConfig), metadata.get("arena", {})) + ) + + env_dict = metadata.get("environment", {}) + if isinstance(env_dict.get("task"), str): + from brittle_star_project.environment.env_types import Task + + try: + env_dict["task"] = Task[env_dict["task"]].name + except Exception: + try: + env_dict["task"] = Task(env_dict["task"]).name + except Exception: + pass + + trained_environment = OmegaConf.to_object( + OmegaConf.merge(OmegaConf.structured(EnvConfig), env_dict) + ) + trained_obs_bounds = OmegaConf.to_object( + OmegaConf.merge( + OmegaConf.structured(ObservationBoundsConfig), metadata.get("obs_bounds", {}) + ) + ) + + return TrainingConfig( + morphology=trained_morphology, + arena=trained_arena, + environment=trained_environment, + obs_bounds=trained_obs_bounds, + ) diff --git a/src/brittle_star_project/evaluation/policy.py b/src/brittle_star_project/evaluation/policy.py new file mode 100644 index 0000000..5aa2d6e --- /dev/null +++ b/src/brittle_star_project/evaluation/policy.py @@ -0,0 +1,89 @@ +from __future__ import annotations + +from pathlib import Path +from typing import Any, Protocol + +import jax +import jax.numpy as jnp +import numpy as np + +from brittle_star_project.evaluation.checkpoint import load_params + + +class ControlPolicy(Protocol): + """Protocol for any policy that can produce actions from observations.""" + + def act(self, *, observations: dict[str, Any]) -> np.ndarray: ... + + +class PolicyAgent: + """Wraps a trained Flax actor for deterministic inference.""" + + def __init__( + self, + *, + sensor_params: Any, + actor_params: Any, + action_dim: int, + obs_processor: Any, + ) -> None: + from brittle_star_project.MLPs.mlps import Actor, GenericDenseLayersWithActivation + + # Infer layer sizes from params + try: + dense_params = ( + sensor_params.get("params", {}) + if isinstance(sensor_params, dict) + else sensor_params["params"] + ) + except Exception: + dense_params = sensor_params + + layer_sizes = [] + idx = 0 + while True: + key = f"Dense_{idx}" + if key not in dense_params: + break + layer_sizes.append(int(np.asarray(dense_params[key]["kernel"]).shape[1])) + idx += 1 + + if not layer_sizes: + raise ValueError("Could not infer Dense_* layers from sensor params") + + self._sensor = GenericDenseLayersWithActivation(layer_sizes=layer_sizes) + self._actor = Actor(action_dim=action_dim) + self._sensor_apply = jax.jit(self._sensor.apply) + self._actor_apply = jax.jit(self._actor.apply) + self._params = { + "sensor_params": sensor_params, + "actor_params": actor_params, + } + self._obs_processor = obs_processor + + @classmethod + def from_checkpoint( + cls, + model_path: Path, + *, + action_dim: int, + obs_processor: Any, + ) -> "PolicyAgent": + """Load params from .flax and construct the agent.""" + params = load_params(model_path) + + return cls( + sensor_params=params["sensor_params"], + actor_params=params["actor_params"], + action_dim=action_dim, + obs_processor=obs_processor, + ) + + def act(self, *, observations: dict[str, Any]) -> np.ndarray: + """Return deterministic action (actor mean, no exploration noise).""" + 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) + + return np.asarray(mean, dtype=np.float32).ravel() diff --git a/src/brittle_star_project/evaluation/rollout.py b/src/brittle_star_project/evaluation/rollout.py new file mode 100644 index 0000000..79c86ed --- /dev/null +++ b/src/brittle_star_project/evaluation/rollout.py @@ -0,0 +1,163 @@ +from __future__ import annotations + +import itertools +import time +from dataclasses import dataclass +from typing import Any + +import numpy as np + +from brittle_star_project import BrittleStarEnv +from brittle_star_project.evaluation.policy import ControlPolicy + + +@dataclass +class EpisodeResult: + return_: float + length: int + reached_target: bool + final_xy_dist: float | None + + +def _get_observations(state: Any) -> dict[str, Any] | None: + return getattr(state, "observations", None) + + +def _get_xy_distance_to_target(observations: dict[str, Any]) -> float | None: + return float(np.asarray(observations["xy_distance_to_target"]).reshape(-1)[0]) + + +def _target_reached(*, state: Any) -> bool: + return bool(getattr(state, "terminated", False) or getattr(state, "truncated", False)) + + +def _maybe_clip_action( + action: np.ndarray, + low: np.ndarray | None, + high: np.ndarray | None, +) -> np.ndarray: + if low is None or high is None: + return action + low = np.asarray(low, dtype=np.float32).ravel() + high = np.asarray(high, dtype=np.float32).ravel() + if low.shape != action.shape or high.shape != action.shape: + return action + return np.clip(action, low, high) + + +def rollout_headless( + *, + env: BrittleStarEnv, + policy: ControlPolicy, + seed: int, + max_steps: int, + action_low: np.ndarray | None, + action_high: np.ndarray | None, + action_mask: np.ndarray | None = None, +) -> EpisodeResult: + """Run an episode headlessly and return the result.""" + state = env.reset(seed=seed) + + ep_return = 0.0 + observations = _get_observations(state) + prev_dist = _get_xy_distance_to_target(observations) if observations else None + reached_target = _target_reached(state=state) + + steps = 0 + for _ in range(int(max_steps)): + obs_dict = observations or {} + + action = policy.act(observations=obs_dict) + if action_mask is not None: + action = action[action_mask] + action = _maybe_clip_action(action, action_low, action_high) + + state = env.step(state=state, action=action) + steps += 1 + + observations = _get_observations(state) + cur_dist = _get_xy_distance_to_target(observations) if observations else None + if prev_dist is not None and cur_dist is not None: + ep_return += prev_dist - cur_dist + prev_dist = cur_dist + + reached_target = _target_reached(state=state) + if reached_target: + break + + final_dist = _get_xy_distance_to_target(observations) if observations else None + return EpisodeResult( + return_=ep_return, + length=steps, + reached_target=reached_target, + final_xy_dist=final_dist, + ) + + +def rollout_viewer( + *, + env: BrittleStarEnv, + policy: ControlPolicy, + seed: int, + state: Any, + control_dt: float, + max_steps: int | None, + action_low: np.ndarray | None, + action_high: np.ndarray | None, + action_mask: np.ndarray | None = None, +) -> None: + """Run an episode using the interactive MuJoCo viewer.""" + import mujoco.viewer + + model = state.mj_model + data = state.mj_data + + episode_return = 0.0 + observations = _get_observations(state) + prev_dist = _get_xy_distance_to_target(observations) if observations else None + reached_target = _target_reached(state=state) + + steps = 0 + with mujoco.viewer.launch_passive(model, data) as viewer: + step_iter = range(int(max_steps)) if max_steps is not None else itertools.count() + for _step_idx in step_iter: + if not viewer.is_running(): + break + step_start = time.time() + + obs_dict = observations or {} + action = policy.act(observations=obs_dict) + if action_mask is not None: + action = action[action_mask] + action = _maybe_clip_action(action, action_low, action_high) + + with viewer.lock(): + state = env.step(state=state, action=action) + + if not viewer.is_running(): + break + viewer.sync() + + steps += 1 + + observations = _get_observations(state) + cur_dist = _get_xy_distance_to_target(observations) if observations else None + if prev_dist is not None and cur_dist is not None: + episode_return += prev_dist - cur_dist + prev_dist = cur_dist + + reached_target = _target_reached(state=state) + if reached_target: + break + + remaining = control_dt - (time.time() - step_start) + if remaining > 0: + time.sleep(remaining) + + dist = _get_xy_distance_to_target(observations) if observations else None + dist_str = "n/a" if dist is None else f"{dist:.3f}" + print( + "episode done: " + f"return={episode_return:.6f}, len={steps}, " + f"target_reached={reached_target}, final_xy_dist={dist_str}" + ) diff --git a/src/brittle_star_project/render/__init__.py b/src/brittle_star_project/render/__init__.py deleted file mode 100644 index 51e0aa1..0000000 --- a/src/brittle_star_project/render/__init__.py +++ /dev/null @@ -1,3 +0,0 @@ -from .renderer import simulate_policy, SimulationConfig, ControlPolicy - -__all__ = ["simulate_policy", "SimulationConfig", "ControlPolicy"] diff --git a/src/brittle_star_project/render/renderer.py b/src/brittle_star_project/render/renderer.py deleted file mode 100644 index 91e669c..0000000 --- a/src/brittle_star_project/render/renderer.py +++ /dev/null @@ -1,78 +0,0 @@ -from __future__ import annotations - -import time -from dataclasses import dataclass -from typing import Any, Protocol - -import numpy as np - - -@dataclass -class SimulationConfig: - realtime: bool = True - seed: int = 0 - - -class ControlPolicy(Protocol): - def act(self, *, obs: np.ndarray | None = None, t: float = 0.0) -> np.ndarray: ... - - -def _default_observations(data: Any) -> np.ndarray: - qpos = np.asarray(data.qpos, dtype=np.float32).ravel() - qvel = np.asarray(data.qvel, dtype=np.float32).ravel() - return np.concatenate([qpos, qvel], axis=0) - - -def simulate_policy( - policy: ControlPolicy, - config: SimulationConfig, - state: Any | None = None, -) -> None: - """Open MuJoCo's native viewer and step using actions from a policy. - - This path drives MuJoCo physics directly (mj_step) and uses the policy output - as `data.ctrl`. - """ - - import mujoco.viewer - - if state is None: - raise ValueError("A valid environment state must be provided.") - - model = state.mj_model - data = state.mj_data - - start = time.time() - with mujoco.viewer.launch_passive(model, data) as viewer: - while viewer.is_running(): - step_start = time.time() - - t = time.time() - start - - # Input vector for the policy - # TODO: custom input - obs = _default_observations(data) - - # Policy action - ctrl = policy.act(obs=obs, t=t) - - # Check if the policy output vector give an input for each actuator (nu) - # TODO: what if model trained on full morphology but we want to test on a damaged one? - # (nu mismatch) - if model.nu > 0: - ctrl = np.asarray(ctrl, dtype=np.float32).ravel() - if ctrl.shape != (model.nu,): - raise ValueError( - f"Policy returned ctrl shape {ctrl.shape}, expected ({model.nu},)" - ) - data.ctrl[:] = ctrl - - # Step the simulation and update the viewer - mujoco.mj_step(model, data) - viewer.sync() - - # If we're running in realtime mode, sleep to maintain real-time pacing. - if config.realtime: - remaining = model.opt.timestep - (time.time() - step_start) - if remaining > 0: - time.sleep(remaining) From d6146850fbae921d08694bbb80cf3a8b7ece3a4f Mon Sep 17 00:00:00 2001 From: Tibo De Peuter Date: Tue, 28 Apr 2026 13:55:36 +0200 Subject: [PATCH 08/18] feat(simulate): capture video --- configs/simulation/default.yaml | 5 + pyproject.toml | 4 + scripts/simulate.py | 48 +++++- .../configs/config_simulation.py | 5 + .../evaluation/__init__.py | 4 + src/brittle_star_project/evaluation/video.py | 148 ++++++++++++++++++ uv.lock | 22 ++- 7 files changed, 234 insertions(+), 2 deletions(-) create mode 100644 src/brittle_star_project/evaluation/video.py diff --git a/configs/simulation/default.yaml b/configs/simulation/default.yaml index 390cd46..b2ca27a 100644 --- a/configs/simulation/default.yaml +++ b/configs/simulation/default.yaml @@ -13,3 +13,8 @@ max_steps: null # Points to a morphology config YAML file (e.g. configs/morphology/3_arms.yaml). # If null, the training morphology from the model's metadata is used. morphology_override: null + +# Video recording (requires [evaluation] extra) +record_video: false +# When null, video is saved in a per-model evaluation folder alongside the model. +video_output_path: null diff --git a/pyproject.toml b/pyproject.toml index 607e704..df3c6a5 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -38,6 +38,10 @@ cuda = [ analysis = [ "tensorboard", ] +evaluation = [ + "imageio>=2.35.0", + "imageio-ffmpeg>=0.5.1", +] [dependency-groups] dev = [ diff --git a/scripts/simulate.py b/scripts/simulate.py index 7bb2800..7f27509 100644 --- a/scripts/simulate.py +++ b/scripts/simulate.py @@ -27,6 +27,11 @@ from brittle_star_project.environment.env_config import MorphologyConfig from brittle_star_project.evaluation.checkpoint import load_metadata, metadata_to_configs from brittle_star_project.evaluation.policy import PolicyAgent from brittle_star_project.evaluation.rollout import rollout_headless, rollout_viewer +from brittle_star_project.evaluation.video import ( + record_episode, + create_evaluation_dir, + save_evaluation_metadata, +) @hydra.main(config_path="../configs", config_name="main_config", version_base="1.3") @@ -118,7 +123,48 @@ def main(dict_cfg: DictConfig) -> None: headless = bool(sim_cfg.headless) max_steps = sim_cfg.max_steps - if headless: + if sim_cfg.record_video: + if max_steps is None: + raise ValueError("simulation.max_steps is required when simulation.record_video=true") + + max_steps_i = int(max_steps) + if max_steps_i <= 0: + raise ValueError("simulation.max_steps must be > 0") + + if sim_cfg.video_output_path is None: + eval_dir = create_evaluation_dir(model_path) + output_path = eval_dir / "simulation.mp4" + else: + output_path = Path(hydra.utils.to_absolute_path(sim_cfg.video_output_path)) + eval_dir = output_path.parent + eval_dir.mkdir(parents=True, exist_ok=True) + + result = record_episode( + env=env, + policy=policy, + seed=seed, + max_steps=max_steps_i, + action_low=action_low, + action_high=action_high, + action_mask=action_mask, + output_path=output_path, + ) + + save_evaluation_metadata( + eval_dir=eval_dir, + morphology_override_path=sim_cfg.morphology_override, + seed=seed, + max_steps=max_steps_i, + result=result, + ) + final_dist_str = "n/a" if result.final_xy_dist is None else f"{result.final_xy_dist:.3f}" + print(f"Video saved to {output_path}") + print( + "episode done: " + f"return={result.return_:.6f}, len={result.length}, " + f"target_reached={result.reached_target}, final_xy_dist={final_dist_str}" + ) + elif headless: if max_steps is None: raise ValueError("simulation.max_steps is required when simulation.headless=true") diff --git a/src/brittle_star_project/configs/config_simulation.py b/src/brittle_star_project/configs/config_simulation.py index a10682d..0b93fd8 100644 --- a/src/brittle_star_project/configs/config_simulation.py +++ b/src/brittle_star_project/configs/config_simulation.py @@ -19,3 +19,8 @@ class SimulationSettings: # Observations are padded from the override morphology UP TO the training # morphology's shape via compute_padding_masks(override, reference=training). morphology_override: Optional[str] = None + + # Video recording (requires [evaluation] extra) + record_video: bool = False + # When None, video is saved in a per-model evaluation folder alongside the model. + video_output_path: Optional[str] = None diff --git a/src/brittle_star_project/evaluation/__init__.py b/src/brittle_star_project/evaluation/__init__.py index 73719e6..38a1c8b 100644 --- a/src/brittle_star_project/evaluation/__init__.py +++ b/src/brittle_star_project/evaluation/__init__.py @@ -3,6 +3,7 @@ from __future__ import annotations from .checkpoint import load_metadata, load_params, metadata_to_configs, TrainingConfig from .policy import PolicyAgent, ControlPolicy from .rollout import rollout_headless, rollout_viewer, EpisodeResult +from .video import record_episode, create_evaluation_dir, save_evaluation_metadata __all__ = [ "load_metadata", @@ -14,4 +15,7 @@ __all__ = [ "rollout_headless", "rollout_viewer", "EpisodeResult", + "record_episode", + "create_evaluation_dir", + "save_evaluation_metadata", ] diff --git a/src/brittle_star_project/evaluation/video.py b/src/brittle_star_project/evaluation/video.py new file mode 100644 index 0000000..726edcb --- /dev/null +++ b/src/brittle_star_project/evaluation/video.py @@ -0,0 +1,148 @@ +from __future__ import annotations + +import datetime +from pathlib import Path + +import numpy as np +import yaml + +from brittle_star_project import BrittleStarEnv +from brittle_star_project.evaluation.policy import ControlPolicy +from brittle_star_project.evaluation.rollout import ( + EpisodeResult, + _get_observations, + _get_xy_distance_to_target, + _target_reached, + _maybe_clip_action, +) + + +def create_evaluation_dir(model_path: Path) -> Path: + """Create a unique timestamped directory for saving evaluation results.""" + timestamp = datetime.datetime.now().strftime("%Y-%m-%d_%H-%M-%S") + eval_dir = model_path.parent / f"{model_path.stem}_evaluations" / f"eval_{timestamp}" + eval_dir.mkdir(parents=True, exist_ok=True) + return eval_dir + + +def save_evaluation_metadata( + eval_dir: Path, + *, + morphology_override_path: str | None, + seed: int, + max_steps: int | None, + result: EpisodeResult, +) -> None: + """Save metadata about the evaluation run.""" + metadata = { + "timestamp": datetime.datetime.now().isoformat(), + "morphology_override": morphology_override_path, + "seed": seed, + "max_steps": max_steps, + "result": { + "return": float(result.return_), + "length": int(result.length), + "reached_target": bool(result.reached_target), + "final_xy_dist": float(result.final_xy_dist) + if result.final_xy_dist is not None + else None, + }, + } + with open(eval_dir / "evaluation_metadata.yaml", "w") as f: + yaml.safe_dump(metadata, f, sort_keys=False) + + +def record_episode( + *, + env: BrittleStarEnv, + policy: ControlPolicy, + seed: int, + max_steps: int, + action_low: np.ndarray | None, + action_high: np.ndarray | None, + action_mask: np.ndarray | None = None, + output_path: Path, + fps: int = 60, + width: int = 640, + height: int = 480, +) -> EpisodeResult: + """Run an episode headlessly and record a video using MuJoCo's Renderer and imageio. + + Args: + env: The environment. + policy: The policy agent. + seed: Random seed. + max_steps: Maximum number of steps. + action_low: Minimum action values. + action_high: Maximum action values. + action_mask: Boolean mask for the actions. + output_path: Where to save the .mp4 file. + fps: Frames per second for the video. + width: Video width. + height: Video height. + """ + try: + import imageio + import mujoco + except ImportError as e: + raise ImportError( + "Video recording requires 'imageio' and 'mujoco'. " + "Please install the evaluation dependencies: `uv pip install .[evaluation]`" + ) from e + + state = env.reset(seed=seed) + model = state.mj_model + data = state.mj_data + + # Use the first camera defined in the environment config, or default to 0 + camera_id = env._config.camera_ids[0] if env._config.camera_ids else 0 + renderer = mujoco.Renderer(model, width=width, height=height) + + ep_return = 0.0 + observations = _get_observations(state) + prev_dist = _get_xy_distance_to_target(observations) if observations else None + reached_target = _target_reached(state=state) + + frames = [] + steps = 0 + + for _ in range(int(max_steps)): + # Capture frame + renderer.update_scene(data, camera=camera_id) + frames.append(renderer.render()) + + # Step environment + obs_dict = observations or {} + action = policy.act(observations=obs_dict) + if action_mask is not None: + action = action[action_mask] + action = _maybe_clip_action(action, action_low, action_high) + + state = env.step(state=state, action=action) + steps += 1 + + observations = _get_observations(state) + cur_dist = _get_xy_distance_to_target(observations) if observations else None + if prev_dist is not None and cur_dist is not None: + ep_return += prev_dist - cur_dist + prev_dist = cur_dist + + reached_target = _target_reached(state=state) + if reached_target: + break + + # Capture final frame + renderer.update_scene(data, camera=camera_id) + frames.append(renderer.render()) + renderer.close() + + # Save video + imageio.mimsave(str(output_path), frames, fps=fps) + + final_dist = _get_xy_distance_to_target(observations) if observations else None + return EpisodeResult( + return_=ep_return, + length=steps, + reached_target=reached_target, + final_xy_dist=final_dist, + ) diff --git a/uv.lock b/uv.lock index bca7772..86d6ed9 100644 --- a/uv.lock +++ b/uv.lock @@ -42,6 +42,10 @@ analysis = [ cuda = [ { name = "jax", extra = ["cuda13"] }, ] +evaluation = [ + { name = "imageio" }, + { name = "imageio-ffmpeg" }, +] [package.dev-dependencies] dev = [ @@ -58,6 +62,8 @@ requires-dist = [ { name = "flax", specifier = ">=0.12.2" }, { name = "gymnasium", specifier = ">=1.2.3" }, { name = "hydra-core", specifier = ">=1.3.2" }, + { name = "imageio", marker = "extra == 'evaluation'", specifier = ">=2.35.0" }, + { name = "imageio-ffmpeg", marker = "extra == 'evaluation'", specifier = ">=0.5.1" }, { name = "ipykernel", specifier = "==7.2.0" }, { name = "jax", specifier = "==0.9.0.1" }, { name = "jax", extras = ["cuda13"], marker = "extra == 'cuda'", specifier = "==0.9.0.1" }, @@ -75,7 +81,7 @@ requires-dist = [ { name = "wandb", specifier = "==0.24.2" }, { name = "warp-lang" }, ] -provides-extras = ["cuda", "analysis"] +provides-extras = ["cuda", "analysis", "evaluation"] [package.metadata.requires-dev] dev = [ @@ -795,6 +801,20 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/49/fa/391e437a34e55095173dca5f24070d89cbc233ff85bf1c29c93248c6588d/imageio-2.37.3-py3-none-any.whl", hash = "sha256:46f5bb8522cd421c0f5ae104d8268f569d856b29eb1a13b92829d1970f32c9f0", size = 317646, upload-time = "2026-03-09T11:31:10.771Z" }, ] +[[package]] +name = "imageio-ffmpeg" +version = "0.6.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/44/bd/c3343c721f2a1b0c9fc71c1aebf1966a3b7f08c2eea8ed5437a2865611d6/imageio_ffmpeg-0.6.0.tar.gz", hash = "sha256:e2556bed8e005564a9f925bb7afa4002d82770d6b08825078b7697ab88ba1755", size = 25210, upload-time = "2025-01-16T21:34:32.747Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/da/58/87ef68ac83f4c7690961bce288fd8e382bc5f1513860fc7f90a9c1c1c6bf/imageio_ffmpeg-0.6.0-py3-none-macosx_10_9_intel.macosx_10_9_x86_64.whl", hash = "sha256:9d2baaf867088508d4a3458e61eeb30e945c4ad8016025545f66c4b5aaef0a61", size = 24932969, upload-time = "2025-01-16T21:34:20.464Z" }, + { url = "https://files.pythonhosted.org/packages/40/5c/f3d8a657d362cc93b81aab8feda487317da5b5d31c0e1fdfd5e986e55d17/imageio_ffmpeg-0.6.0-py3-none-macosx_11_0_arm64.whl", hash = "sha256:b1ae3173414b5fc5f538a726c4e48ea97edc0d2cdc11f103afee655c463fa742", size = 21113891, upload-time = "2025-01-16T21:34:00.277Z" }, + { url = "https://files.pythonhosted.org/packages/33/e7/1925bfbc563c39c1d2e82501d8372734a5c725e53ac3b31b4c2d081e895b/imageio_ffmpeg-0.6.0-py3-none-manylinux2014_aarch64.whl", hash = "sha256:1d47bebd83d2c5fc770720d211855f208af8a596c82d17730aa51e815cdee6dc", size = 25632706, upload-time = "2025-01-16T21:33:53.475Z" }, + { url = "https://files.pythonhosted.org/packages/a0/2d/43c8522a2038e9d0e7dbdf3a61195ecc31ca576fb1527a528c877e87d973/imageio_ffmpeg-0.6.0-py3-none-manylinux2014_x86_64.whl", hash = "sha256:c7e46fcec401dd990405049d2e2f475e2b397779df2519b544b8aab515195282", size = 29498237, upload-time = "2025-01-16T21:34:13.726Z" }, + { url = "https://files.pythonhosted.org/packages/a0/13/59da54728351883c3c1d9fca1710ab8eee82c7beba585df8f25ca925f08f/imageio_ffmpeg-0.6.0-py3-none-win32.whl", hash = "sha256:196faa79366b4a82f95c0f4053191d2013f4714a715780f0ad2a68ff37483cc2", size = 19652251, upload-time = "2025-01-16T21:34:06.812Z" }, + { url = "https://files.pythonhosted.org/packages/2c/c6/fa760e12a2483469e2bf5058c5faff664acf66cadb4df2ad6205b016a73d/imageio_ffmpeg-0.6.0-py3-none-win_amd64.whl", hash = "sha256:02fa47c83703c37df6bfe4896aab339013f62bf02c5ebf2dce6da56af04ffc0a", size = 31246824, upload-time = "2025-01-16T21:34:28.6Z" }, +] + [[package]] name = "iniconfig" version = "2.3.0" From aabf1bc1862ca798c7daea50eff3e069834993b0 Mon Sep 17 00:00:00 2001 From: Tibo De Peuter Date: Tue, 28 Apr 2026 14:00:12 +0200 Subject: [PATCH 09/18] feat(simulate): configurable camera id --- configs/simulation/default.yaml | 2 ++ scripts/simulate.py | 1 + src/brittle_star_project/configs/config_simulation.py | 2 ++ src/brittle_star_project/evaluation/video.py | 4 ++-- 4 files changed, 7 insertions(+), 2 deletions(-) diff --git a/configs/simulation/default.yaml b/configs/simulation/default.yaml index b2ca27a..e689c3b 100644 --- a/configs/simulation/default.yaml +++ b/configs/simulation/default.yaml @@ -18,3 +18,5 @@ morphology_override: null record_video: false # When null, video is saved in a per-model evaluation folder alongside the model. video_output_path: null +# Camera ID to use for video recording (1 is usually the close-up camera) +camera_id: 1 diff --git a/scripts/simulate.py b/scripts/simulate.py index 7f27509..f17d8ef 100644 --- a/scripts/simulate.py +++ b/scripts/simulate.py @@ -148,6 +148,7 @@ def main(dict_cfg: DictConfig) -> None: action_high=action_high, action_mask=action_mask, output_path=output_path, + camera_id=sim_cfg.camera_id, ) save_evaluation_metadata( diff --git a/src/brittle_star_project/configs/config_simulation.py b/src/brittle_star_project/configs/config_simulation.py index 0b93fd8..e15d94e 100644 --- a/src/brittle_star_project/configs/config_simulation.py +++ b/src/brittle_star_project/configs/config_simulation.py @@ -24,3 +24,5 @@ class SimulationSettings: record_video: bool = False # When None, video is saved in a per-model evaluation folder alongside the model. video_output_path: Optional[str] = None + # Camera ID to use for video recording (1 is usually the close-up camera) + camera_id: int = 1 diff --git a/src/brittle_star_project/evaluation/video.py b/src/brittle_star_project/evaluation/video.py index 726edcb..174326b 100644 --- a/src/brittle_star_project/evaluation/video.py +++ b/src/brittle_star_project/evaluation/video.py @@ -62,6 +62,7 @@ def record_episode( action_high: np.ndarray | None, action_mask: np.ndarray | None = None, output_path: Path, + camera_id: int = 1, fps: int = 60, width: int = 640, height: int = 480, @@ -77,6 +78,7 @@ def record_episode( action_high: Maximum action values. action_mask: Boolean mask for the actions. output_path: Where to save the .mp4 file. + camera_id: Camera index to use for rendering (1 is usually close-up). fps: Frames per second for the video. width: Video width. height: Video height. @@ -94,8 +96,6 @@ def record_episode( model = state.mj_model data = state.mj_data - # Use the first camera defined in the environment config, or default to 0 - camera_id = env._config.camera_ids[0] if env._config.camera_ids else 0 renderer = mujoco.Renderer(model, width=width, height=height) ep_return = 0.0 From a1b4df0921f59b92d61c1f72f624bb8c5ed3ed80 Mon Sep 17 00:00:00 2001 From: Tibo De Peuter Date: Tue, 28 Apr 2026 14:20:27 +0200 Subject: [PATCH 10/18] test(simulate): evaluation --- tests/test_evaluation.py | 77 ++++++++++++++++++++++++++++++++++++++++ 1 file changed, 77 insertions(+) create mode 100644 tests/test_evaluation.py diff --git a/tests/test_evaluation.py b/tests/test_evaluation.py new file mode 100644 index 0000000..c8bffa9 --- /dev/null +++ b/tests/test_evaluation.py @@ -0,0 +1,77 @@ +import numpy as np + +from brittle_star_project.evaluation.checkpoint import metadata_to_configs, TrainingConfig +from brittle_star_project.evaluation.rollout import _maybe_clip_action +from brittle_star_project.environment.env_config import ( + MorphologyConfig, + ArenaConfig, + EnvConfig, + ObservationBoundsConfig, +) +from brittle_star_project.environment.env_types import Task + + +def test_metadata_to_configs(): + """Test that a raw metadata dictionary correctly instantiates the typed configs.""" + mock_metadata = { + "morphology": { + "segments_per_arm": [4, 0, 4, 0, 0], + "use_p_control": False, + }, + "arena": {"sand_ground_color": False, "size": [15.0, 10.0]}, + "environment": { + "task": "LIGHT_ESCAPE", + "simulation_time": 5000.0, + }, + "obs_bounds": {"joint_velocity": [-10.0, 10.0]}, + } + + config = metadata_to_configs(mock_metadata) + + assert isinstance(config, TrainingConfig) + + # Check MorphologyConfig + assert isinstance(config.morphology, MorphologyConfig) + assert config.morphology.segments_per_arm == [4, 0, 4, 0, 0] + assert config.morphology.use_p_control is False + assert config.morphology.use_torque_control is False # default + + # Check ArenaConfig + assert isinstance(config.arena, ArenaConfig) + assert config.arena.sand_ground_color is False + assert config.arena.size == [15.0, 10.0] + assert config.arena.wall_height == 1.5 # default + + # Check EnvConfig + assert isinstance(config.environment, EnvConfig) + assert config.environment.task == Task.LIGHT_ESCAPE + assert config.environment.simulation_time == 5000.0 + assert config.environment.time_scale == 2 # default + + # Check ObservationBoundsConfig + assert isinstance(config.obs_bounds, ObservationBoundsConfig) + assert config.obs_bounds.joint_velocity == [-10.0, 10.0] + assert config.obs_bounds.segment_contact == [0.0, 1.0] # default + + +def test_maybe_clip_action(): + """Test action clipping against boundaries.""" + # Test valid clipping + action = np.array([1.5, -2.5, 0.0]) + low = np.array([-1.0, -1.0, -1.0]) + high = np.array([1.0, 1.0, 1.0]) + + clipped = _maybe_clip_action(action, low, high) + np.testing.assert_array_equal(clipped, np.array([1.0, -1.0, 0.0])) + + # Test skipping when bounds are None + unclipped_1 = _maybe_clip_action(action, None, high) + np.testing.assert_array_equal(unclipped_1, action) + + unclipped_2 = _maybe_clip_action(action, low, None) + np.testing.assert_array_equal(unclipped_2, action) + + # Test skipping on shape mismatch + wrong_low = np.array([-1.0, -1.0]) # Shape mismatch + unclipped_3 = _maybe_clip_action(action, wrong_low, high) + np.testing.assert_array_equal(unclipped_3, action) From 5716b73eb9620b4092d6910bd003ac000fb4079a Mon Sep 17 00:00:00 2001 From: Tibo De Peuter Date: Tue, 28 Apr 2026 16:25:05 +0200 Subject: [PATCH 11/18] docs: restructure API docs --- README.md | 44 ++++++++----------------------- docs/DEVELOPMENT.md | 21 +++------------ docs/README.md | 10 +++++-- docs/api/simulate.md | 14 ---------- docs/api/simulation.md | 39 +++++++++++++++++++++++++++ docs/api/tracking.md | 60 ++++++++++++++++++++++++++++++++++++++++++ docs/api/training.md | 49 ++++++++++++++++++++++++++++++++++ 7 files changed, 171 insertions(+), 66 deletions(-) delete mode 100644 docs/api/simulate.md create mode 100644 docs/api/simulation.md create mode 100644 docs/api/tracking.md create mode 100644 docs/api/training.md diff --git a/README.md b/README.md index dd85e12..572fd73 100644 --- a/README.md +++ b/README.md @@ -13,44 +13,22 @@ To set up the UV module, you can run the following command: uv sync --frozen ``` -### Configuration +## Usage -1. **Copy the default configuration:** +For detailed instructions on how to use the project, please refer to the **[API Documentation](docs/README.md)**. + +### Quick Start + +1. **Train a model:** ```bash - cp configs/default_ppo.yaml configs/my_experiment.yaml + uv run python scripts/train.py ppo.learning_rate=0.001 logging.track=true ``` -2. **Edit `configs/my_experiment.yaml`** to set your WandB credentials: - ```yaml - track: true # Enable WandB logging - wandb_entity: "your-wandb-username" # Replace with your username/team - wandb_project_name: "PPO-Modularity" - ``` +2. **Monitor progress:** + See [Tracking & Monitoring](docs/api/tracking.md). -3. **(Optional) Login to WandB:** - ```bash - uv run wandb login - ``` - -### Training - -example command: - -```bash -uv run python scripts/train.py -``` - -Or use a custom config file: - -```bash -uv run python scripts/train.py --config configs/my_experiment.yaml -``` - -Override specific parameters: - -```bash -uv run python scripts/train.py --learning-rate 0.001 --num-envs 32 --track -``` +3. **Simulate a trained model:** + See [Simulation & Evaluation](docs/api/simulation.md). ## HPC diff --git a/docs/DEVELOPMENT.md b/docs/DEVELOPMENT.md index 5a01d26..3518b38 100644 --- a/docs/DEVELOPMENT.md +++ b/docs/DEVELOPMENT.md @@ -62,22 +62,9 @@ In the devcontainer, this will succeed on both CPU and GPU. A `GpuDevice` is exp ## Logging & Monitoring -This project uses a unified logging system through the `experiment_logger` package. For a full API reference, see the [package README](../src/experiment_logger/README.md). - -### Quick Setup - -1. **Authorization**: Export your API key in your terminal to enable WandB synchronization: - ```bash - export WANDB_API_KEY=your_copied_api_key_here - ``` -2. **Toggle Tracking**: Use the `--track` flag in `scripts/train.py` to enable online sync. -3. **Local Monitoring**: All runs are recorded in the `runs/` directory. View scalars with TensorBoard: - ```bash - tensorboard --logdir runs/ - ``` - -### Environment Awareness - -The logger automatically detects if it is running in an interactive terminal or a non-interactive environment (like an HPC Slurm job). It will automatically disable progress bars and switch to robust fallback modes (offline logging) to ensure your experiments never hang. +This project uses a unified logging system through the `experiment_logger` package. +- **Usage in Code**: To use the logger in your scripts, refer to the [package README](../src/experiment_logger/README.md) for the API reference. +- **WandB/TensorBoard Setup**: For information on how to configure tracking for experiments, see the [Tracking & Monitoring API Guide](./api/tracking.md). +The logger automatically detects if it is running in an interactive terminal or a non-interactive environment (like an HPC Slurm job), adjusting progress bars and fallback modes accordingly. diff --git a/docs/README.md b/docs/README.md index a4f4ea0..3c76684 100644 --- a/docs/README.md +++ b/docs/README.md @@ -2,6 +2,8 @@ ## Design & architecture ([`/design`](./design/)) +If you are interested in the "why did you do it like this?" + - [Actor/critic architecture](./design/actor-critic.md): Description of the actor-critic pipeline. - [Communication](./design/communication.md): Message propagation, Nerve-Net style. - [Controllers](./design/controllers.md): Macroscopig brain toplogy, centralized, arm-level, segment-level. @@ -11,5 +13,9 @@ ## API reference ([`/api`](./api/)) -- [Environment](./api/environment.md): MuJoCo environment interaction, state retrieval, and configuration. -- [Simulate](./api/simulate.md): Simulation rendering. +If you are interested in the "how do I use it?" + +- [Training](./api/training.md): How to configure and run experiments. +- [Tracking & Monitoring](./api/tracking.md): Setting up WandB and TensorBoard to monitor runs. +- [Simulation](./api/simulation.md): Visualizing and evaluating models. +- [Environment](./api/environment.md): MuJoCo environment interaction and configuration. diff --git a/docs/api/simulate.md b/docs/api/simulate.md deleted file mode 100644 index 9c4b21b..0000000 --- a/docs/api/simulate.md +++ /dev/null @@ -1,14 +0,0 @@ -# Training and Simulation for Brittle Star Models - -## Simulating a model - -In order to simulate and view the behavior of a trained model, you can use the `simulate.py` script. This script allows you to specify the path to a trained model and will launch a simulation using that model. This script has the following parameters: - -- `--model`: The path to the trained model artifact to simulate. -- `--model-type`: The type of model to simulate (e.g., `random`, ...) -- `--task`: The task to simulate (e.g., `directed_locomotion`, ...) -- `--seed`: The random seed for reproducibility. - -```bash -python simulate.py --model artifacts/my_model --model-type random --task directed_locomotion --seed 0 -``` \ No newline at end of file diff --git a/docs/api/simulation.md b/docs/api/simulation.md new file mode 100644 index 0000000..0a0e041 --- /dev/null +++ b/docs/api/simulation.md @@ -0,0 +1,39 @@ +# Simulation & Evaluation + +The simulation pipeline allows you to visualize trained models and evaluate their performance under various conditions. + +## Overview + +The simulation pipeline is metadata-driven. Training-specific configurations (morphology, arena, environment, etc.) are automatically loaded from the `_metadata.yaml` file associated with the model checkpoint. + +## Basic Simulation + +To simulate a model in the MuJoCo viewer: + +```bash +uv run scripts/simulate.py simulation.model_path=runs/your_run/final_model.flax +``` + +## Amputation & Morphology Overrides + +You can test trained models on different morphologies (e.g., amputating legs) by providing a morphology override. The observations will be automatically padded up to the training morphology's dimensions: + +```bash +uv run scripts/simulate.py \ + simulation.model_path=runs/your_run/final_model.flax \ + simulation.morphology_override=configs/morphology/3_arms.yaml +``` + +## Video Recording + +Recording videos requires the `[evaluation]` extra: + +```bash +uv run scripts/simulate.py \ + simulation.model_path=runs/your_run/final_model.flax \ + simulation.record_video=true \ + simulation.max_steps=1000 +``` + +Videos and evaluation metadata are stored in timestamped folders alongside the model: +`runs/your_run/final_model_evaluations/eval_/simulation.mp4` diff --git a/docs/api/tracking.md b/docs/api/tracking.md new file mode 100644 index 0000000..226992b --- /dev/null +++ b/docs/api/tracking.md @@ -0,0 +1,60 @@ +# Tracking & Monitoring + +This guide explains how to monitor your experiments using Weights & Biases (WandB) and TensorBoard. + +## Weights & Biases (WandB) + +WandB is used for online synchronization and visualization of training metrics. + +### Authorization + +Export your API key in your terminal to enable WandB synchronization: + +```bash +export WANDB_API_KEY=your_copied_api_key_here +``` + +Alternatively, you can log in using the CLI: + +```bash +uv run wandb login +``` + +### Enabling Tracking + +To enable online sync during a training run, set `logging.track=true` on the command line: + +```bash +uv run python scripts/train.py logging.track=true +``` + +You can also configure your project and entity: + +```bash +uv run python scripts/train.py \ + logging.track=true \ + logging.wandb_project_name="MyProject" \ + logging.wandb_entity="my-team" +``` + +These can also be set in your configuration YAML file under the `logging` key. + +## Local Monitoring with TensorBoard + +All runs are recorded locally in the `runs/` directory (or the directory specified in `experiment.base_run_dir`). You can view scalars and other metrics with TensorBoard: + +```bash +tensorboard --logdir runs/ +``` + +Access the interface at `http://localhost:6006`. + +### CLI Exploration Tool + +For quick diagnostics or to export data to CSV without launching the full TensorBoard UI, you can use the `explore_tensorboard.py` script: + +```bash +uv run python scripts/analysis/explore_tensorboard.py runs/your_run_name/ +``` + +See the detailed description in [`/scripts/analysis/README.md`](../../scripts/analysis/README.md). diff --git a/docs/api/training.md b/docs/api/training.md new file mode 100644 index 0000000..a9678bb --- /dev/null +++ b/docs/api/training.md @@ -0,0 +1,49 @@ +# Training Models + +This guide covers how to configure and run training experiments for the Brittle Star project using Hydra-based configurations. + +## Configuration + +The project uses a modular configuration system powered by [Hydra](https://hydra.cc/). Instead of passing many command-line flags, you select and override configuration groups. + +### Creating a Custom Experiment + +1. **Create a new experiment file:** + Create a file at `configs/experiment/my_experiment.yaml`. You can copy an existing one as a template: + ```bash + cp configs/experiment/base.yaml configs/experiment/my_experiment.yaml + ``` + +2. **Edit `configs/experiment/my_experiment.yaml`** to set your experiment parameters: + ```yaml + # @package _global_ + experiment: + exp_name: "my_custom_run" + seed: 42 + ``` + +## Training Execution + +To start a training run with the default settings defined in `configs/main_config.yaml`: + +```bash +uv run python scripts/train.py +``` + +### Using a Custom Experiment Configuration + +To run with your custom experiment file: + +```bash +uv run python scripts/train.py experiment=my_experiment +``` + +### Command-Line Overrides + +You can override any parameter directly from the command line using Hydra's dot notation. This is useful for quick tests: + +```bash +uv run python scripts/train.py ppo.learning_rate=0.001 ppo.num_envs=32 logging.track=true +``` + +For more details on tracking your experiments, see [Tracking & Monitoring](./tracking.md). From 33fbc4c96f6f564f2282f0aedde3a08c95153d42 Mon Sep 17 00:00:00 2001 From: Tibo De Peuter Date: Tue, 28 Apr 2026 16:42:46 +0200 Subject: [PATCH 12/18] feat: tool to construct XML --- scripts/analysis/dump_mjcf.py | 186 ++++++++++++++++++++++++++++++++ tests/test_morphology_render.py | 4 +- 2 files changed, 188 insertions(+), 2 deletions(-) create mode 100644 scripts/analysis/dump_mjcf.py diff --git a/scripts/analysis/dump_mjcf.py b/scripts/analysis/dump_mjcf.py new file mode 100644 index 0000000..1f552d9 --- /dev/null +++ b/scripts/analysis/dump_mjcf.py @@ -0,0 +1,186 @@ +#!/usr/bin/env python3 +""" +Dump MJCF XML for a brittle-star morphology using the project's Hydra configs. + +Usage examples: + + # Use a named morphology config from configs/morphology (Hydra style) + uv run python scripts/analysis/dump_mjcf.py morphology=3_arms + + # Use a morphology override YAML (same key as simulation.morphology_override) + uv run python scripts/analysis/dump_mjcf.py \ + simulation.morphology_override=configs/morphology/3_arms.yaml + +Output path: + Provide `dump_out=path/to/file.xml` on the command line, otherwise writes `morphology.xml` in + current directory. +""" + +from __future__ import annotations + +from pathlib import Path +import sys +import xml.etree.ElementTree as ET + +import hydra +from omegaconf import DictConfig, OmegaConf + +from brittle_star_project.configs.main_config import BrittleStarConfig +from brittle_star_project.configs.register_configs import register_configs +from brittle_star_project.environment.env_config import MorphologyConfig +from brittle_star_project.environment.factory import BrittleStarEnvFactory + + +def try_serialize(obj): + """Try multiple common accessors to obtain an XML string from the morphology object.""" + candidates = [ + "to_xml_string", + "to_xml", + "to_string", + "to_mjcf", + "to_mjcf_string", + "get_mjcf", + "get_mjcf_str", + "get_mjcf_assets", + "export_to_xml_with_assets", + "get_xml", + "xml", + "mjcf", + "mjcf_model", + "mjcf_body", + "model", + "root", + ] + + def normalize(out): + if out is None: + return None + # lxml element + try: + import lxml.etree as lxml_et + + if isinstance(out, lxml_et._Element): + return lxml_et.tostring(out, encoding="unicode") + except Exception: + pass + + if isinstance(out, ET.Element): + return ET.tostring(out, encoding="unicode") + + if isinstance(out, bytes): + try: + return out.decode() + except Exception: + return None + + if hasattr(out, "toxml") and callable(out.toxml): + try: + return out.toxml() + except Exception: + pass + + try: + s = str(out) + if s.lstrip().startswith("<"): + return s + return s + except Exception: + return None + + for name in candidates: + attr = getattr(obj, name, None) + if callable(attr): + try: + out = attr() + except Exception: + out = None + if out: + norm = normalize(out) + if norm: + return norm + elif attr is not None: + norm = normalize(attr) + if norm: + return norm + + if hasattr(obj, "mjcf"): + nested = getattr(obj, "mjcf") + if nested is not None: + return try_serialize(nested) + + return None + + +@hydra.main(config_path="../../configs", config_name="main_config", version_base="1.3") +def main(dict_cfg: DictConfig) -> None: + # Compose typed config like the rest of the project + cfg = OmegaConf.to_object(OmegaConf.merge(OmegaConf.structured(BrittleStarConfig), dict_cfg)) + + # Check for a simulation morphology override (points to a YAML file) + sim_override = None + try: + sim_override = dict_cfg.get("simulation", {}).get("morphology_override", None) + except Exception: + sim_override = getattr(getattr(cfg, "simulation", None), "morphology_override", None) + + if sim_override: + override_path = Path(hydra.utils.to_absolute_path(sim_override)) + if not override_path.exists(): + raise FileNotFoundError(f"Could not find morphology override YAML at {override_path}") + import yaml + + with open(override_path, "r") as f: + override_dict = yaml.safe_load(f) + env_morphology = OmegaConf.to_object( + OmegaConf.merge(OmegaConf.structured(MorphologyConfig), override_dict) + ) + else: + env_morphology = cfg.morphology + + # Ensure we have a MorphologyConfig instance + if isinstance(env_morphology, dict): + morph_cfg = MorphologyConfig(**env_morphology) + else: + morph_cfg = env_morphology + + # Build morphology via project factory (same as runtime) + morph = BrittleStarEnvFactory.create_morphology(morph_cfg) + + xml_text = try_serialize(morph) + if xml_text is None and hasattr(morph, "mjcf"): + xml_text = try_serialize(morph.mjcf) + + if xml_text is None: + raise RuntimeError( + "Failed to serialize morphology to MJCF/XML. Inspect the `morph` object interactively." + ) + + # Prefer explicit CLI override `dump_out=...` if provided, otherwise choose a sensible default. + dump_out = None + try: + dump_out = dict_cfg.get("dump_out", None) + except Exception: + dump_out = None + + if dump_out is None: + # If the user passed a morphology group on the CLI (e.g. morphology=3_arms), + # use a descriptive default path under `runs/morphologies/`. + morph_name = None + for a in sys.argv[1:]: + if a.startswith("morphology="): + morph_name = a.split("=", 1)[1] + break + + default_out = f"runs/morphologies/{morph_name}.xml" if morph_name else "morphology.xml" + out_path = Path(hydra.utils.to_absolute_path(default_out)) + else: + out_path = Path(hydra.utils.to_absolute_path(str(dump_out))) + + out_path.parent.mkdir(parents=True, exist_ok=True) + out_path.write_text(xml_text) + print(f"Wrote MJCF XML to {out_path}") + + +if __name__ == "__main__": + register_configs() + main() diff --git a/tests/test_morphology_render.py b/tests/test_morphology_render.py index 6b435cf..a80b5e6 100644 --- a/tests/test_morphology_render.py +++ b/tests/test_morphology_render.py @@ -14,7 +14,7 @@ from brittle_star_project.environment.BrittleStarJaxEnvWrapper import BrittleSta @pytest.mark.skipif(os.getenv("CI") == "true", reason="No OpenGL display in CI") def test_render_morphologies(): - base_dir = "runs/renders" + base_dir = "runs/morphologies" os.makedirs(base_dir, exist_ok=True) # --- 1. Full 5-Arm Morphology --- @@ -34,7 +34,7 @@ def test_render_morphologies(): renderer_full = mujoco.Renderer(model=model_full) renderer_full.update_scene(data_full, camera=1) pixels_full = renderer_full.render() - image_path = os.path.join(base_dir, "full_5_arm.png") + image_path = os.path.join(base_dir, "5_arm.png") Image.fromarray(pixels_full).save(image_path) print(f"Generated full morphology render: {image_path}") From 49f587403524e3474b2dcb4c491a1e4e12fe2386 Mon Sep 17 00:00:00 2001 From: Tibo De Peuter Date: Tue, 28 Apr 2026 17:27:08 +0200 Subject: [PATCH 13/18] chore: remove distance to target from model input --- docs/design/input_action_spaces.md | 7 +++++-- src/brittle_star_project/environment/env_config.py | 2 -- src/brittle_star_project/environment/obs_processing.py | 1 - 3 files changed, 5 insertions(+), 5 deletions(-) diff --git a/docs/design/input_action_spaces.md b/docs/design/input_action_spaces.md index 143dce5..50f142e 100644 --- a/docs/design/input_action_spaces.md +++ b/docs/design/input_action_spaces.md @@ -16,8 +16,7 @@ Global inputs, always broadcasted to all nodes: $$ tilt = sqrt(roll^2 + pitch^2) $$ -- Goal vector: Instead of just a scalar distance, the goal is represented asa a vector (distance and ange/direction) to - the target. +- Goal vector: Instead of just a scalar distance, the goal is represented as an angle/direction to the target. Local inputs, routed directly to specific nodes: @@ -70,6 +69,10 @@ When designing the state space, we must ask: *Could a human operator perform thi as a normalized unit vector bounds the values to the $[-1, 1]$ range, which stabilizes neural network training. Providing only a scalar "distance to the goal" would force the agent to learning localized searching behaviors (e.g. random walks or spiraling) to deduce the direction, drastically increasing the difficulty of the learning task. + + **NOTE:** We later dropped the "distance to vector", switching to only a direction as the input. Our reasoning is + the agent should always move towards the goal (it should not learn to stop at the goal), which allows for this + simplification that decreases the model input size. - Contact sensors: Segment contact detects external ground interaction and is biologically vital for timing gait transitions. - **Zero-Centered Rescaling ($[-1, 1]$):** Using a zero-centered range is standard best practice for continuous control diff --git a/src/brittle_star_project/environment/env_config.py b/src/brittle_star_project/environment/env_config.py index 0af387e..b9a621d 100644 --- a/src/brittle_star_project/environment/env_config.py +++ b/src/brittle_star_project/environment/env_config.py @@ -72,7 +72,6 @@ class ObservationBoundsConfig: 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]]: @@ -82,6 +81,5 @@ class ObservationBoundsConfig: "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 index 0ef6e3d..37a775f 100644 --- a/src/brittle_star_project/environment/obs_processing.py +++ b/src/brittle_star_project/environment/obs_processing.py @@ -63,7 +63,6 @@ def create_obs_processor( "joint_velocity", "segment_contact", "unit_xy_direction_to_target", - "xy_distance_to_target", ] values = [] for key in ordered_keys: From 2a34b84f83125c2503554ab9bbb83f3cfd16a9e4 Mon Sep 17 00:00:00 2001 From: Tibo De Peuter Date: Tue, 28 Apr 2026 17:38:02 +0200 Subject: [PATCH 14/18] chore: min/max of inputs --- scripts/extract_observation_bounds.py | 138 ++++++++++++++++++ .../environment/env_config.py | 13 +- 2 files changed, 147 insertions(+), 4 deletions(-) create mode 100644 scripts/extract_observation_bounds.py diff --git a/scripts/extract_observation_bounds.py b/scripts/extract_observation_bounds.py new file mode 100644 index 0000000..374babc --- /dev/null +++ b/scripts/extract_observation_bounds.py @@ -0,0 +1,138 @@ +#!/usr/bin/env python3 +"""Empirically extract observation bounds (focused on joint velocities). + +This script creates a MuJoCo environment using the project's factory and +randomly samples actions to discover observed maxima for selected +observation keys (joint_velocity, joint_position, joint_actuator_force). + +Usage: + python scripts/extract_observation_bounds.py \ + --morphology configs/morphology/3_arms.yaml --num-steps 5000 --seed 42 + +If `--morphology` is omitted the default `MorphologyConfig()` is used. +""" + +from __future__ import annotations + +import argparse +from pathlib import Path +import yaml +import numpy as np + +from brittle_star_project import BrittleStarEnvFactory, BrittleStarEnv, Backend +from brittle_star_project.environment.env_config import ( + MorphologyConfig, + ArenaConfig, + EnvConfig, +) + + +def load_morphology(path: str | None) -> MorphologyConfig: + if path is None: + return MorphologyConfig() + p = Path(path) + if not p.exists(): + raise FileNotFoundError(f"Morphology file not found: {p}") + with open(p, "r") as f: + data = yaml.safe_load(f) or {} + return MorphologyConfig(**data) + + +def _extract_observations(state): + # Under different backends the returned state may be a dict or an object + obs = getattr(state, "observations", None) + if obs is None and isinstance(state, dict): + obs = state.get("observations", state) + return obs + + +def find_empirical_bounds( + morph_cfg: MorphologyConfig, + arena_cfg: ArenaConfig, + env_cfg: EnvConfig, + num_steps: int = 5000, + seed: int = 42, +) -> None: + factory = BrittleStarEnvFactory() + raw_env = factory.create_environment(Backend.MJC, morph_cfg, arena_cfg, env_cfg) + env = BrittleStarEnv(raw_env, backend=Backend.MJC, config=env_cfg, morphology_config=morph_cfg) + + # Initial reset + state = env.reset(seed=seed) + + # Determine action bounds + action_space = getattr(raw_env, "action_space", None) + if action_space is None: + raise RuntimeError("Environment missing `action_space`; cannot sample actions.") + + action_low = np.asarray(action_space.low, dtype=np.float32) + action_high = np.asarray(action_space.high, dtype=np.float32) + action_shape = action_low.shape + + # Track maximum absolute observed values + tracked_keys = ["joint_velocity", "joint_position", "joint_actuator_force"] + max_observed = {k: 0.0 for k in tracked_keys} + + # Include observation at reset + obs0 = _extract_observations(state) + if isinstance(obs0, dict): + for k in tracked_keys: + if k in obs0: + max_observed[k] = max(max_observed[k], float(np.max(np.abs(np.asarray(obs0[k]))))) + + rng = np.random.RandomState(seed) + for i in range(num_steps): + u = rng.uniform(size=action_shape) + action = action_low + (action_high - action_low) * u + + # Provide a numpy RNG to the env step; wrapper will pass it if accepted. + step_out = env.step(state=state, action=action, rng=env.make_rng(seed + i + 1)) + + # Unpack next state from common return conventions + if hasattr(step_out, "state"): + next_state = step_out.state + elif isinstance(step_out, (tuple, list)) and len(step_out) >= 1: + next_state = step_out[0] + else: + next_state = step_out + + obs = _extract_observations(next_state) + if isinstance(obs, dict): + for k in tracked_keys: + if k in obs: + val = float(np.max(np.abs(np.asarray(obs[k])))) + if val > max_observed[k]: + max_observed[k] = val + + state = next_state + + # Print recommended bounds with a 20% safety margin + print("\n--- Recommended Observation Bounds (20% margin) ---") + for k, v in max_observed.items(): + if v == 0.0: + print(f"{k}: observed max 0.0 (increase sampling or inspect env)") + else: + safe = v * 1.2 + print(f"{k}: [-{safe:.6f}, {safe:.6f}] (observed max: {v:.6f})") + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument( + "--morphology", type=str, default=None, help="Path to morphology YAML (optional)" + ) + parser.add_argument( + "--num-steps", type=int, default=5000, help="Number of random steps to sample" + ) + parser.add_argument("--seed", type=int, default=42, help="RNG seed") + args = parser.parse_args() + + morph_cfg = load_morphology(args.morphology) + arena_cfg = ArenaConfig() + env_cfg = EnvConfig() + + find_empirical_bounds(morph_cfg, arena_cfg, env_cfg, num_steps=args.num_steps, seed=args.seed) + + +if __name__ == "__main__": + main() diff --git a/src/brittle_star_project/environment/env_config.py b/src/brittle_star_project/environment/env_config.py index b9a621d..f26b428 100644 --- a/src/brittle_star_project/environment/env_config.py +++ b/src/brittle_star_project/environment/env_config.py @@ -66,10 +66,15 @@ class EnvConfig: 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]) + # Empirical testing based on the extract_observation_bounds.py script run for 1.000.000 steps + + # Based on max. ctrlrange (0.78539816339744828) in XML, but empirical testing went slightly over + joint_position: list[float] = field(default_factory=lambda: [-0.8, 0.8]) + # Empirical testing showed max. 3.22, adding buffer to be safe + joint_velocity: list[float] = field(default_factory=lambda: [-5.0, 5.0]) + # Based on max. forceRange in XML, verified with empirical testing + joint_actuator_force: list[float] = field(default_factory=lambda: [-3.75, 3.75]) + # Based on intuition and reasoning 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]) disk_z_tilt: list[float] = field(default_factory=lambda: [0.0, 3.141592653589793]) From ba8184b03423a0e2e4825224dd035d41b29235a5 Mon Sep 17 00:00:00 2001 From: Tibo De Peuter Date: Tue, 28 Apr 2026 18:03:53 +0200 Subject: [PATCH 15/18] test: fix outdated imports etc. --- tests/test_network_shapes.py | 19 ++++--------------- 1 file changed, 4 insertions(+), 15 deletions(-) diff --git a/tests/test_network_shapes.py b/tests/test_network_shapes.py index 32e9898..e1a7a1a 100644 --- a/tests/test_network_shapes.py +++ b/tests/test_network_shapes.py @@ -2,8 +2,8 @@ import jax import jax.numpy as jnp from brittle_star_project.environment.padded_obs_wrapper import ( compute_padding_masks, - pad_observations_batched, ) +from brittle_star_project.environment.obs_processing import create_obs_processor # We use Actor and OneDenseLayerMLP (as the critic) based on your mlps.py from brittle_star_project.MLPs.mlps import Actor, OneDenseLayerMLP @@ -20,21 +20,10 @@ def test_centralized_forward_pass_with_padding(): "segment_contact": jnp.zeros((batch_size, 14)), } - # 2. Pad Observation using the boolean scattering wrapper + # 2. Process and Pad Observation masks = compute_padding_masks(segments_per_arm=(4, 0, 4, 2, 4)) - padded_obs = pad_observations_batched(amputated_obs, masks) - - # Assertions to ensure padding sizes are correct (40 joints, 20 segments) - assert padded_obs["joint_position"].shape == (batch_size, 40), "Padding failed for joint keys" - assert padded_obs["segment_contact"].shape == (batch_size, 20), ( - "Padding failed for segment keys" - ) - - # 3. Concatenate for Centralized MLP (simulating the global state vector) - global_state = jnp.concatenate( - [padded_obs["joint_position"], padded_obs["joint_velocity"], padded_obs["segment_contact"]], - axis=-1, - ) + obs_processor = create_obs_processor(bounds_dict={}, padding_masks=masks) + global_state = obs_processor(amputated_obs) # 40 + 40 + 20 = 100 dimensions assert global_state.shape == (batch_size, 100), ( From be78fb15bd6e111a2f2d7ef91bac122aa06e0647 Mon Sep 17 00:00:00 2001 From: Tibo De Peuter Date: Thu, 30 Apr 2026 20:19:42 +0200 Subject: [PATCH 16/18] feat: allow model metadata cli override --- configs/simulation/default.yaml | 4 ++ scripts/simulate.py | 6 ++- .../configs/config_simulation.py | 4 ++ .../evaluation/checkpoint.py | 12 +++--- tests/test_evaluation.py | 39 ++++++++++++++++++- 5 files changed, 58 insertions(+), 7 deletions(-) diff --git a/configs/simulation/default.yaml b/configs/simulation/default.yaml index e689c3b..61599a4 100644 --- a/configs/simulation/default.yaml +++ b/configs/simulation/default.yaml @@ -20,3 +20,7 @@ record_video: false video_output_path: null # Camera ID to use for video recording (1 is usually the close-up camera) camera_id: 1 + +# Optional override for the metadata YAML file path. +# If null, the script looks for `_metadata.yaml` alongside the model_path. +metadata_path: null diff --git a/scripts/simulate.py b/scripts/simulate.py index f17d8ef..a2b121e 100644 --- a/scripts/simulate.py +++ b/scripts/simulate.py @@ -51,7 +51,11 @@ def main(dict_cfg: DictConfig) -> None: raise ValueError(f"Expected a '.flax' checkpoint, got '{model_path.name}'.") # 2. Discover + load sidecar metadata YAML - metadata = load_metadata(model_path) + metadata_override = None + if sim_cfg.metadata_path is not None: + metadata_override = Path(hydra.utils.to_absolute_path(sim_cfg.metadata_path)) + + metadata = load_metadata(model_path, metadata_override) # 3. Reconstruct typed configs from metadata training = metadata_to_configs(metadata) diff --git a/src/brittle_star_project/configs/config_simulation.py b/src/brittle_star_project/configs/config_simulation.py index e15d94e..9fd1507 100644 --- a/src/brittle_star_project/configs/config_simulation.py +++ b/src/brittle_star_project/configs/config_simulation.py @@ -26,3 +26,7 @@ class SimulationSettings: video_output_path: Optional[str] = None # Camera ID to use for video recording (1 is usually the close-up camera) camera_id: int = 1 + + # Optional override for the sidecar metadata YAML file. + # If None, it defaults to the model_path with a `_metadata.yaml` suffix. + metadata_path: Optional[str] = None diff --git a/src/brittle_star_project/evaluation/checkpoint.py b/src/brittle_star_project/evaluation/checkpoint.py index ffbc27a..d10046b 100644 --- a/src/brittle_star_project/evaluation/checkpoint.py +++ b/src/brittle_star_project/evaluation/checkpoint.py @@ -56,13 +56,15 @@ def load_params(path: Path) -> dict: } -def load_metadata(model_path: Path) -> dict: +def load_metadata(model_path: Path, metadata_override_path: Path | None = None) -> dict: """Discover and load the sidecar metadata YAML file.""" - metadata_path = model_path.with_name(model_path.stem + "_metadata.yaml") + if metadata_override_path is not None: + metadata_path = metadata_override_path + else: + metadata_path = model_path.with_name(model_path.stem + "_metadata.yaml") + if not metadata_path.exists(): - raise FileNotFoundError( - f"Could not find metadata YAML for {model_path.name}. Expected it at {metadata_path}" - ) + raise FileNotFoundError(f"Could not find metadata YAML at {metadata_path}") with open(metadata_path, "r") as f: return yaml.safe_load(f) diff --git a/tests/test_evaluation.py b/tests/test_evaluation.py index c8bffa9..bbb03f8 100644 --- a/tests/test_evaluation.py +++ b/tests/test_evaluation.py @@ -1,6 +1,13 @@ import numpy as np +import pytest +import yaml +from pathlib import Path -from brittle_star_project.evaluation.checkpoint import metadata_to_configs, TrainingConfig +from brittle_star_project.evaluation.checkpoint import ( + metadata_to_configs, + TrainingConfig, + load_metadata, +) from brittle_star_project.evaluation.rollout import _maybe_clip_action from brittle_star_project.environment.env_config import ( MorphologyConfig, @@ -75,3 +82,33 @@ def test_maybe_clip_action(): wrong_low = np.array([-1.0, -1.0]) # Shape mismatch unclipped_3 = _maybe_clip_action(action, wrong_low, high) np.testing.assert_array_equal(unclipped_3, action) + + +def test_load_metadata_with_override(tmp_path: Path): + """Test that metadata can be loaded from both default and override paths.""" + # 1. Setup + model_path = tmp_path / "model.flax" + model_path.write_bytes(b"dummy") + + default_metadata_path = tmp_path / "model_metadata.yaml" + default_content = {"version": "default", "seed": 42} + with open(default_metadata_path, "w") as f: + yaml.dump(default_content, f) + + override_path = tmp_path / "custom_metadata.yaml" + override_content = {"version": "override", "seed": 1337} + with open(override_path, "w") as f: + yaml.dump(override_content, f) + + # 2. Test default behavior + loaded_default = load_metadata(model_path) + assert loaded_default == default_content + + # 3. Test override behavior + loaded_override = load_metadata(model_path, metadata_override_path=override_path) + assert loaded_override == override_content + + # 4. Test Error Case + non_existent = tmp_path / "missing.yaml" + with pytest.raises(FileNotFoundError, match="Could not find metadata YAML at"): + load_metadata(model_path, metadata_override_path=non_existent) From 86a53ee9f7b2ecea8122451ba96cda0d27c367c9 Mon Sep 17 00:00:00 2001 From: Tibo De Peuter Date: Thu, 30 Apr 2026 21:10:17 +0200 Subject: [PATCH 17/18] fix: robot direction to target --- docs/design/input_action_spaces.md | 28 +++---- .../environment/env_config.py | 10 +-- .../environment/obs_processing.py | 11 ++- tests/test_target_direction.py | 78 +++++++++++++++++++ 4 files changed, 107 insertions(+), 20 deletions(-) create mode 100644 tests/test_target_direction.py diff --git a/docs/design/input_action_spaces.md b/docs/design/input_action_spaces.md index 50f142e..ddb7a55 100644 --- a/docs/design/input_action_spaces.md +++ b/docs/design/input_action_spaces.md @@ -16,7 +16,8 @@ Global inputs, always broadcasted to all nodes: $$ tilt = sqrt(roll^2 + pitch^2) $$ -- Goal vector: Instead of just a scalar distance, the goal is represented as an angle/direction to the target. +- Goal vector: A 2D unit vector representing the *egocentric* direction to the target. A value of $[1.0, 0.0]$ + indicates that the target is directly in front of the agent (angle 0). Local inputs, routed directly to specific nodes: @@ -73,20 +74,21 @@ When designing the state space, we must ask: *Could a human operator perform thi **NOTE:** We later dropped the "distance to vector", switching to only a direction as the input. Our reasoning is the agent should always move towards the goal (it should not learn to stop at the goal), which allows for this simplification that decreases the model input size. + + The environment provides a raw `unit_xy_direction_to_target` (global), which we transform into a calculated + `robot_direction_to_target` (egocentric) before passing it to the MLPs. This vector consists of the X and Y + direction, where a value of $[1.0, 0.0]$ (mapping to an angle of $0$) means the robot is facing directly towards the + target. - Contact sensors: Segment contact detects external ground interaction and is biologically vital for timing gait transitions. -- **Zero-Centered Rescaling ($[-1, 1]$):** Using a zero-centered range is standard best practice for continuous control +- Zero-Centered Rescaling ($[-1, 1]$): Using a zero-centered range is standard best practice for continuous control tasks. It provides several mathematical and physical advantages: - - **Improved Gradient Flow:** Neural networks optimize faster when inputs are zero-centered. If all inputs were - positive (e.g., $[0, 1]$), the gradients during backpropagation would be forced to the same sign, causing - inefficient "zig-zag" weight updates. - - **Meaningful "Neutral" State:** In robotics, $0.0$ naturally represents a resting state (zero velocity, centered + - Improved Gradient Flow: Neural networks optimize faster when inputs are zero-centered. If all inputs were positive + (e.g., $[0, 1]$), the gradients during backpropagation would be forced to the same sign, causing inefficient + "zig-zag" weight updates. + - Meaningful Neutral State: In robotics, $0.0$ naturally represents a resting state (zero velocity, centered position, no force). In a $[-1, 1]$ system, this physical rest maps to a neutral $0.0$ signal in the network. - - **Robustness to Amputation:** In this project, amputated limbs are padded with $0.0$. In a $[-1, 1]$ system, this - correctly communicates a "neutral/dead" signal. In a $[0, 1]$ system, $0.0$ would represent the absolute minimum - physical limit, causing the network to misinterpret missing limbs as being at their extreme limits. - - + This also correctly communicaties a "neutral/dead" signal for amputated limbs that are padded with $0.0$ values. Specifically, we do not include some available inputs: @@ -120,8 +122,7 @@ This is what the filtered input vectors look like in MuJoCo, with $J$ joints and - `joint_velocity`: shape=(J,), dtype=float64 - `joint_actuator_force`: shape=(J,), dtype=float64 - `segment_contact`: shape=(S,), dtype=float64 -- `unit_xy_direction_to_target`: shape=(2,), dtype=float64 -- `xy_distance_to_target`: shape=(1,), dtype=float64 +- `robot_direction_to_target`: shape=(2,), dtype=float64, egocentric - `disk_z_tilt`: shape=(1,), dtype=float64, derived from `disk_rotation` This brings the entire input space down to $3J + S + 4$ float64's, compared to $4J + S + 15$ float64's for the @@ -159,6 +160,5 @@ disk_angular_velocity: shape=(3,), dtype=float64, size=3 tendon_position: shape=(0,), dtype=float64, size=0 tendon_velocity: shape=(0,), dtype=float64, size=0 segment_contact: shape=(6,), dtype=float64, size=6 -unit_xy_direction_to_target: shape=(2,), dtype=float64, size=2 xy_distance_to_target: shape=(1,), dtype=float64, size=1 ``` diff --git a/src/brittle_star_project/environment/env_config.py b/src/brittle_star_project/environment/env_config.py index f26b428..22b8967 100644 --- a/src/brittle_star_project/environment/env_config.py +++ b/src/brittle_star_project/environment/env_config.py @@ -70,21 +70,21 @@ class ObservationBoundsConfig: # Based on max. ctrlrange (0.78539816339744828) in XML, but empirical testing went slightly over joint_position: list[float] = field(default_factory=lambda: [-0.8, 0.8]) - # Empirical testing showed max. 3.22, adding buffer to be safe + # Empirical testing showed max. 3.22, adding buffer to be safe. Consider higher values "fast". joint_velocity: list[float] = field(default_factory=lambda: [-5.0, 5.0]) # Based on max. forceRange in XML, verified with empirical testing joint_actuator_force: list[float] = field(default_factory=lambda: [-3.75, 3.75]) # Based on intuition and reasoning 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]) + robot_direction_to_target: list[float] = field(default_factory=lambda: [-1.0, 1.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 { + "disk_z_tilt": tuple(self.disk_z_tilt), + "joint_actuator_force": tuple(self.joint_actuator_force), "joint_position": tuple(self.joint_position), "joint_velocity": tuple(self.joint_velocity), - "joint_actuator_force": tuple(self.joint_actuator_force), + "robot_direction_to_target": tuple(self.robot_direction_to_target), "segment_contact": tuple(self.segment_contact), - "unit_xy_direction_to_target": tuple(self.unit_xy_direction_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 index 37a775f..b18ee8c 100644 --- a/src/brittle_star_project/environment/obs_processing.py +++ b/src/brittle_star_project/environment/obs_processing.py @@ -26,6 +26,15 @@ def create_obs_processor( 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)) + + if "unit_xy_direction_to_target" in new_obs: + yaw = rot[2] + unit_x, unit_y = new_obs["unit_xy_direction_to_target"] + cos_yaw, sin_yaw = jnp.cos(yaw), jnp.sin(yaw) + new_x = unit_x * cos_yaw + unit_y * sin_yaw + new_y = -unit_x * sin_yaw + unit_y * cos_yaw + new_obs["robot_direction_to_target"] = jnp.stack([new_x, new_y]) + return new_obs def _normalize_features(obs: dict) -> dict: @@ -61,8 +70,8 @@ def create_obs_processor( "joint_actuator_force", "joint_position", "joint_velocity", + "robot_direction_to_target", "segment_contact", - "unit_xy_direction_to_target", ] values = [] for key in ordered_keys: diff --git a/tests/test_target_direction.py b/tests/test_target_direction.py new file mode 100644 index 0000000..a81cdf1 --- /dev/null +++ b/tests/test_target_direction.py @@ -0,0 +1,78 @@ +import jax.numpy as jnp + +from brittle_star_project.configs.main_config import BrittleStarConfig +from brittle_star_project.environment.env_types import Backend +from brittle_star_project.environment.BrittleStarJaxEnvWrapper import BrittleStarJaxEnvWrapper +from brittle_star_project.environment.obs_processing import create_obs_processor + + +def test_raw_environment_returns_allocentric_direction(): + """ + Verifies that the raw environment returns a GLOBAL (allocentric) + direction to the target. If the robot rotates in place, + the global vector to the target should remain identical. + """ + env = BrittleStarJaxEnvWrapper.default(num_envs=1, backend=Backend.MJX) + env_state = env.reset(seed=42) + raw_obs_1 = env_state.observations["unit_xy_direction_to_target"] + + ninety_deg_z_quat = jnp.array([0.7071068, 0.0, 0.0, 0.7071068]) + new_qpos = env_state.mjx_data.qpos.at[..., 3:7].set(ninety_deg_z_quat) + new_data = env_state.mjx_data.replace(qpos=new_qpos) + rotated_env_state = env_state.replace(mjx_data=new_data) + + zero_action = jnp.zeros(env.single_action_space.shape) + if len(raw_obs_1.shape) > 1: + zero_action = jnp.expand_dims(zero_action, 0) + final_env_state = env.step(rotated_env_state, zero_action) + raw_obs_2 = final_env_state.observations["unit_xy_direction_to_target"] + + # If the vector is allocentric, it should not change when the robot spins. + assert jnp.sum(jnp.abs(raw_obs_1 - raw_obs_2)) < 1e-4, ( + f"The raw environment observation changed when the robot rotated! " + f"This means it is already egocentric. " + f"Obs 1: {raw_obs_1}, Obs 2: {raw_obs_2}" + ) + + +def test_processor_converts_to_egocentric_direction(): + """ + Verifies that the obs_processor correctly applies a 2D inverse rotation + matrix to convert the global target vector into a local (egocentric) vector. + """ + cfg = BrittleStarConfig() + env = BrittleStarJaxEnvWrapper.default(num_envs=1, backend=Backend.MJX) + + obs_processor = create_obs_processor( + bounds_dict=cfg.obs_bounds.to_bounds_dict(), padding_masks=env.padding_masks + ) + + env_state = env.reset(seed=42) + + # --- Scenario 1 --- + # Robot is rotated 90 degrees Left (facing global Y) + # Target is straight ahead on the global X axis [1.0, 0.0] + # Because the robot is facing Y, the target on X is to its RIGHT [0.0, -1.0] locally. + dummy_obs_1 = dict(env_state.observations) + dummy_obs_1["disk_rotation"] = jnp.array([[0.0, 0.0, jnp.pi / 2.0]]) + dummy_obs_1["unit_xy_direction_to_target"] = jnp.array([[1.0, 0.0]]) + processed_1 = obs_processor(dummy_obs_1) + + # --- Scenario 2 (used to find the array indices) --- + # We change ONLY the target vector so we can isolate it in the final array + dummy_obs_2 = dict(env_state.observations) + dummy_obs_2["disk_rotation"] = jnp.array([[0.0, 0.0, jnp.pi / 2.0]]) + dummy_obs_2["unit_xy_direction_to_target"] = jnp.array([[0.0, 1.0]]) + processed_2 = obs_processor(dummy_obs_2) + + # Find the indices of the elements that changed + diff_array = jnp.abs(processed_1[0] - processed_2[0]) + changed_indices = jnp.where(diff_array > 1e-4)[0] + + local_target = processed_1[0, changed_indices] + expected_local_target = jnp.array([0.0, -1.0]) + + assert jnp.sum(jnp.abs(local_target - expected_local_target)) < 1e-4, ( + f"The obs_processor did not correctly rotate the vector to egocentric. " + f"Expected {expected_local_target}, but got {local_target}." + ) From 304a8e9c4336224e5b50d15b61967524ba525ce7 Mon Sep 17 00:00:00 2001 From: Tibo De Peuter Date: Thu, 30 Apr 2026 22:18:07 +0200 Subject: [PATCH 18/18] refactor: clean tools --- scripts/analysis/dump_mjcf.py | 186 ------------------ scripts/tools/dump_mjcf.py | 141 +++++++++++++ .../{ => tools}/extract_observation_bounds.py | 0 3 files changed, 141 insertions(+), 186 deletions(-) delete mode 100644 scripts/analysis/dump_mjcf.py create mode 100644 scripts/tools/dump_mjcf.py rename scripts/{ => tools}/extract_observation_bounds.py (100%) diff --git a/scripts/analysis/dump_mjcf.py b/scripts/analysis/dump_mjcf.py deleted file mode 100644 index 1f552d9..0000000 --- a/scripts/analysis/dump_mjcf.py +++ /dev/null @@ -1,186 +0,0 @@ -#!/usr/bin/env python3 -""" -Dump MJCF XML for a brittle-star morphology using the project's Hydra configs. - -Usage examples: - - # Use a named morphology config from configs/morphology (Hydra style) - uv run python scripts/analysis/dump_mjcf.py morphology=3_arms - - # Use a morphology override YAML (same key as simulation.morphology_override) - uv run python scripts/analysis/dump_mjcf.py \ - simulation.morphology_override=configs/morphology/3_arms.yaml - -Output path: - Provide `dump_out=path/to/file.xml` on the command line, otherwise writes `morphology.xml` in - current directory. -""" - -from __future__ import annotations - -from pathlib import Path -import sys -import xml.etree.ElementTree as ET - -import hydra -from omegaconf import DictConfig, OmegaConf - -from brittle_star_project.configs.main_config import BrittleStarConfig -from brittle_star_project.configs.register_configs import register_configs -from brittle_star_project.environment.env_config import MorphologyConfig -from brittle_star_project.environment.factory import BrittleStarEnvFactory - - -def try_serialize(obj): - """Try multiple common accessors to obtain an XML string from the morphology object.""" - candidates = [ - "to_xml_string", - "to_xml", - "to_string", - "to_mjcf", - "to_mjcf_string", - "get_mjcf", - "get_mjcf_str", - "get_mjcf_assets", - "export_to_xml_with_assets", - "get_xml", - "xml", - "mjcf", - "mjcf_model", - "mjcf_body", - "model", - "root", - ] - - def normalize(out): - if out is None: - return None - # lxml element - try: - import lxml.etree as lxml_et - - if isinstance(out, lxml_et._Element): - return lxml_et.tostring(out, encoding="unicode") - except Exception: - pass - - if isinstance(out, ET.Element): - return ET.tostring(out, encoding="unicode") - - if isinstance(out, bytes): - try: - return out.decode() - except Exception: - return None - - if hasattr(out, "toxml") and callable(out.toxml): - try: - return out.toxml() - except Exception: - pass - - try: - s = str(out) - if s.lstrip().startswith("<"): - return s - return s - except Exception: - return None - - for name in candidates: - attr = getattr(obj, name, None) - if callable(attr): - try: - out = attr() - except Exception: - out = None - if out: - norm = normalize(out) - if norm: - return norm - elif attr is not None: - norm = normalize(attr) - if norm: - return norm - - if hasattr(obj, "mjcf"): - nested = getattr(obj, "mjcf") - if nested is not None: - return try_serialize(nested) - - return None - - -@hydra.main(config_path="../../configs", config_name="main_config", version_base="1.3") -def main(dict_cfg: DictConfig) -> None: - # Compose typed config like the rest of the project - cfg = OmegaConf.to_object(OmegaConf.merge(OmegaConf.structured(BrittleStarConfig), dict_cfg)) - - # Check for a simulation morphology override (points to a YAML file) - sim_override = None - try: - sim_override = dict_cfg.get("simulation", {}).get("morphology_override", None) - except Exception: - sim_override = getattr(getattr(cfg, "simulation", None), "morphology_override", None) - - if sim_override: - override_path = Path(hydra.utils.to_absolute_path(sim_override)) - if not override_path.exists(): - raise FileNotFoundError(f"Could not find morphology override YAML at {override_path}") - import yaml - - with open(override_path, "r") as f: - override_dict = yaml.safe_load(f) - env_morphology = OmegaConf.to_object( - OmegaConf.merge(OmegaConf.structured(MorphologyConfig), override_dict) - ) - else: - env_morphology = cfg.morphology - - # Ensure we have a MorphologyConfig instance - if isinstance(env_morphology, dict): - morph_cfg = MorphologyConfig(**env_morphology) - else: - morph_cfg = env_morphology - - # Build morphology via project factory (same as runtime) - morph = BrittleStarEnvFactory.create_morphology(morph_cfg) - - xml_text = try_serialize(morph) - if xml_text is None and hasattr(morph, "mjcf"): - xml_text = try_serialize(morph.mjcf) - - if xml_text is None: - raise RuntimeError( - "Failed to serialize morphology to MJCF/XML. Inspect the `morph` object interactively." - ) - - # Prefer explicit CLI override `dump_out=...` if provided, otherwise choose a sensible default. - dump_out = None - try: - dump_out = dict_cfg.get("dump_out", None) - except Exception: - dump_out = None - - if dump_out is None: - # If the user passed a morphology group on the CLI (e.g. morphology=3_arms), - # use a descriptive default path under `runs/morphologies/`. - morph_name = None - for a in sys.argv[1:]: - if a.startswith("morphology="): - morph_name = a.split("=", 1)[1] - break - - default_out = f"runs/morphologies/{morph_name}.xml" if morph_name else "morphology.xml" - out_path = Path(hydra.utils.to_absolute_path(default_out)) - else: - out_path = Path(hydra.utils.to_absolute_path(str(dump_out))) - - out_path.parent.mkdir(parents=True, exist_ok=True) - out_path.write_text(xml_text) - print(f"Wrote MJCF XML to {out_path}") - - -if __name__ == "__main__": - register_configs() - main() diff --git a/scripts/tools/dump_mjcf.py b/scripts/tools/dump_mjcf.py new file mode 100644 index 0000000..ae582c7 --- /dev/null +++ b/scripts/tools/dump_mjcf.py @@ -0,0 +1,141 @@ +#!/usr/bin/env python3 +""" +Dump MJCF XML for a brittle-star morphology using the project's Hydra configs. + +Usage examples: + + # Use a named morphology config from configs/morphology (Hydra style) + uv run python scripts/analysis/dump_mjcf.py morphology=3_arms + + # Use a morphology override YAML (same key as simulation.morphology_override) + uv run python scripts/analysis/dump_mjcf.py \ + simulation.morphology_override=configs/morphology/3_arms.yaml + +Output path: + Provide `dump_out=path/to/file.xml` on the command line, otherwise writes `morphology.xml` in + current directory or `runs/morphologies/.xml`. +""" + +from __future__ import annotations + +import dataclasses +import logging +import sys +from pathlib import Path +from typing import Any, Optional + +import hydra +import yaml +from omegaconf import DictConfig, OmegaConf + +from brittle_star_project.configs.register_configs import register_configs +from brittle_star_project.environment.env_config import MorphologyConfig +from brittle_star_project.environment.factory import BrittleStarEnvFactory + +logging.basicConfig(level=logging.INFO, format="%(levelname)s: %(message)s") +logger = logging.getLogger(__name__) + + +def extract_xml_string(obj: Any) -> Optional[str]: + """ + Attempts to serialize the morphology object to an XML string by checking + common dm_control and internal API methods. + """ + serialization_methods = [ + "to_xml_string", + "to_xml", + "to_string", + "to_mjcf", + "to_mjcf_string", + "get_mjcf", + "get_mjcf_str", + "export_to_xml_string", + ] + + # If the object itself has an 'mjcf' attribute, try to serialize that instead + target_obj = getattr(obj, "mjcf", obj) + + for method_name in serialization_methods: + method = getattr(target_obj, method_name, None) + if callable(method): + try: + xml_data = method() + # Safely handle both string and byte responses + if isinstance(xml_data, str): + return xml_data + elif isinstance(xml_data, bytes): + return xml_data.decode("utf-8") + except Exception as e: + logger.debug(f"Method {method_name}() failed during serialization: {e}") + + return None + + +def resolve_output_path(cfg: DictConfig) -> Path: + """Determines the appropriate output path for the MJCF XML.""" + dump_out = cfg.get("dump_out", None) + if dump_out is not None: + return Path(hydra.utils.to_absolute_path(str(dump_out))) + + morph_name = "morphology" + for arg in sys.argv[1:]: + if arg.startswith("morphology="): + morph_name = arg.split("=", 1)[1] + break + + default_out = ( + f"runs/morphologies/{morph_name}.xml" if morph_name != "morphology" else "morphology.xml" + ) + return Path(hydra.utils.to_absolute_path(default_out)) + + +@hydra.main(config_path="../../configs", config_name="main_config", version_base="1.3") +def main(cfg: DictConfig) -> None: + """Main entry point to construct the morphology and dump its XML.""" + logger.info("Initializing morphology construction...") + + # Extract morphology config safely using dict `.get()` to avoid OmegaConf AttributeErrors + simulation_cfg = cfg.get("simulation", cfg) + override_path = simulation_cfg.get("morphology_override", None) + + if override_path: + logger.info(f"Using morphology override: {override_path}") + with open(hydra.utils.to_absolute_path(override_path), "r") as f: + data = yaml.safe_load(f) or {} + morph_cfg = MorphologyConfig(**data) + else: + # Fallback to default simulation morphology, or an empty base config + morph_node = simulation_cfg.get("morphology", cfg.get("morphology", None)) + + if morph_node is not None: + # Convert OmegaConf node to dict and instantiate MorphologyConfig. + # This ensures any missing keys gracefully fall back to the dataclass defaults. + morph_dict = OmegaConf.to_container(morph_node, resolve=True) + if isinstance(morph_dict, dict): + # Filter to avoid unexpected kwargs if the dataclass is strictly defined + if dataclasses.is_dataclass(MorphologyConfig): + valid_keys = {f.name for f in dataclasses.fields(MorphologyConfig)} + morph_dict = {k: v for k, v in morph_dict.items() if k in valid_keys} + morph_cfg = MorphologyConfig(**morph_dict) + else: + morph_cfg = MorphologyConfig() + else: + morph_cfg = MorphologyConfig() + + morphology = BrittleStarEnvFactory.create_morphology(morph_cfg) + + xml_text = extract_xml_string(morphology) + if not xml_text: + raise RuntimeError("Failed to serialize morphology to MJCF/XML. ") + + out_path = resolve_output_path(cfg) + out_path.parent.mkdir(parents=True, exist_ok=True) + with out_path.open("w", encoding="utf-8") as f: + f.write(xml_text) + + logger.info(f"Successfully exported MJCF XML to: {out_path}") + + +if __name__ == "__main__": + register_configs() + main() diff --git a/scripts/extract_observation_bounds.py b/scripts/tools/extract_observation_bounds.py similarity index 100% rename from scripts/extract_observation_bounds.py rename to scripts/tools/extract_observation_bounds.py