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/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/configs/simulation/default.yaml b/configs/simulation/default.yaml index 84694e0..61599a4 100644 --- a/configs/simulation/default.yaml +++ b/configs/simulation/default.yaml @@ -9,7 +9,18 @@ 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 + +# 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 +# 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/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). diff --git a/docs/design/input_action_spaces.md b/docs/design/input_action_spaces.md index 7860e16..ddb7a55 100644 --- a/docs/design/input_action_spaces.md +++ b/docs/design/input_action_spaces.md @@ -5,12 +5,26 @@ 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 - the target. +Global inputs, always broadcasted to all nodes: + +- 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: 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: + +- 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)** @@ -18,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?* @@ -26,15 +50,50 @@ 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. + + **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 + 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. + 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: + +- 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 +109,56 @@ 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 + +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 +- `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 +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 +xy_distance_to_target: shape=(1,), dtype=float64, size=1 +``` 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 a9a51f4..a2b121e 100644 --- a/scripts/simulate.py +++ b/scripts/simulate.py @@ -1,600 +1,118 @@ """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=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 +from omegaconf import DictConfig, OmegaConf import yaml -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, +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 + +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, ) -_ALLOWED_OBS_KEYS = { - "joint_position", - "joint_velocity", - "joint_actuator_force", - "actuator_force", - "disk_position", - "disk_rotation", - "disk_linear_velocity", - "disk_angular_velocity", - "unit_xy_direction_to_target", - "xy_distance_to_target", -} - - -def _dense_layer_sizes_from_params(params: Any) -> list[int]: - """Infer GenericDenseLayersWithActivation.layer_sizes from a Flax params tree.""" - - 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) - - -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__( - self, - *, - sensor_params: Any, - actor_params: Any, - action_dim: int, - ) -> None: - from brittle_star_project.MLPs.mlps import Actor, GenericDenseLayersWithActivation - - layer_sizes = _dense_layer_sizes_from_params(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, - } - - @staticmethod - def load( - path: Path, - *, - action_dim: int, - ) -> "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." - ) - - 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." - ) - - return CleanRLPPOPolicy( - sensor_params=sensor_params, - actor_params=actor_params, - action_dim=action_dim, - ) - - def act(self, *, observations: dict[str, Any]) -> np.ndarray: - obs = _transform_obs_dict(observations) - 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 _rollout_one_episode_headless( - *, - env: BrittleStarEnv, - policy: CleanRLPPOPolicy, - seed: int, - 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. - - 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 - 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 {} - 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) - - 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 - - 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 _run_one_episode_viewer( - *, - env: BrittleStarEnv, - policy: CleanRLPPOPolicy, - seed: int, - state: Any, - control_dt: float, - max_steps: int | None, - action_low: np.ndarray | None, - action_high: np.ndarray | None, - padding_masks: dict[str, Any] | 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) - 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 {} - 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) - 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(): - 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 _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 - ) - - 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) - @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_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) + + # 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 = 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=training.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, + training.arena, + training.environment, ) env = BrittleStarEnv( raw_env, backend=backend, - config=config.environment, - morphology_config=config.morphology, + config=training.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(training.morphology.segments_per_arm) * 2 + + # 7. Load policy + policy = PolicyAgent.from_checkpoint( + model_path, action_dim=trained_action_dim, obs_processor=obs_processor + ) + + # 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) @@ -605,63 +123,87 @@ def main(dict_cfg: DictConfig) -> None: None if action_space is None else np.asarray(action_space.high, dtype=np.float32).ravel() ) - # ======= MODEL SETUP ======= - nu = int(state0.mj_model.nu) - policy = CleanRLPPOPolicy.load(model_path, action_dim=nu) + # 8. Run simulation + headless = bool(sim_cfg.headless) + max_steps = sim_cfg.max_steps - # 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]) - 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." - ) - - # ======= SIMULATION ======= - headless = bool(config.simulation.headless) - max_steps = config.simulation.max_steps - - if headless: + if sim_cfg.record_video: if max_steps is None: - raise ValueError("simulation.max_steps is required when simulation.headless=true") + 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") - ep_return, ep_len, reached_target, final_dist = _rollout_one_episode_headless( + 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, - padding_masks=padding_masks, + action_mask=action_mask, + output_path=output_path, + camera_id=sim_cfg.camera_id, ) - final_dist_str = "n/a" if final_dist is None else f"{final_dist:.3f}" + + 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={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}" + ) + elif 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") + + result = 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 result.final_xy_dist is None else f"{result.final_xy_dist:.3f}" + print( + "episode done: " + 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 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(training.environment.num_physics_steps_per_control_step) - _run_one_episode_viewer( + rollout_viewer( env=env, policy=policy, seed=seed, @@ -670,7 +212,7 @@ def main(dict_cfg: DictConfig) -> None: max_steps=max_steps_val, action_low=action_low, action_high=action_high, - padding_masks=padding_masks, + action_mask=action_mask, ) env.close() 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/tools/extract_observation_bounds.py b/scripts/tools/extract_observation_bounds.py new file mode 100644 index 0000000..374babc --- /dev/null +++ b/scripts/tools/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/__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/configs/config_simulation.py b/src/brittle_star_project/configs/config_simulation.py index 872ca20..9fd1507 100644 --- a/src/brittle_star_project/configs/config_simulation.py +++ b/src/brittle_star_project/configs/config_simulation.py @@ -13,7 +13,20 @@ 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 + + # 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 + # 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/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/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/environment/__init__.py b/src/brittle_star_project/environment/__init__.py index 4395b1f..7a9a9eb 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, MorphMode 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,7 +12,8 @@ __all__ = [ "Backend", "Task", "BrittleStarEnv", - "StepResult", "BrittleStarEnvFactory", "MorphMode", + "create_obs_processor", + "compute_padding_masks", ] diff --git a/src/brittle_star_project/environment/env_config.py b/src/brittle_star_project/environment/env_config.py index 6683e45..adf0736 100644 --- a/src/brittle_star_project/environment/env_config.py +++ b/src/brittle_star_project/environment/env_config.py @@ -70,4 +70,30 @@ class EnvConfig: # Per docs in upstream env config: integer factors of 200. light_perlin_noise_scale: int = 0 - morph_mode: MorphMode = MorphMode.CENTRALIZED + +@dataclass +class ObservationBoundsConfig: + """Physical observation bounds for deterministic min-max normalization.""" + + # 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. 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]) + 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), + "robot_direction_to_target": tuple(self.robot_direction_to_target), + "segment_contact": tuple(self.segment_contact), + } 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..b18ee8c --- /dev/null +++ b/src/brittle_star_project/environment/obs_processing.py @@ -0,0 +1,91 @@ +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)) + + 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: + 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", + "robot_direction_to_target", + "segment_contact", + ] + 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)) 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 diff --git a/src/brittle_star_project/evaluation/__init__.py b/src/brittle_star_project/evaluation/__init__.py new file mode 100644 index 0000000..38a1c8b --- /dev/null +++ b/src/brittle_star_project/evaluation/__init__.py @@ -0,0 +1,21 @@ +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", + "load_params", + "metadata_to_configs", + "TrainingConfig", + "PolicyAgent", + "ControlPolicy", + "rollout_headless", + "rollout_viewer", + "EpisodeResult", + "record_episode", + "create_evaluation_dir", + "save_evaluation_metadata", +] diff --git a/src/brittle_star_project/evaluation/checkpoint.py b/src/brittle_star_project/evaluation/checkpoint.py new file mode 100644 index 0000000..d10046b --- /dev/null +++ b/src/brittle_star_project/evaluation/checkpoint.py @@ -0,0 +1,107 @@ +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, metadata_override_path: Path | None = None) -> dict: + """Discover and load the sidecar metadata YAML file.""" + 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 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/evaluation/video.py b/src/brittle_star_project/evaluation/video.py new file mode 100644 index 0000000..174326b --- /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, + camera_id: int = 1, + 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. + 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. + """ + 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 + + 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/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) diff --git a/src/brittle_star_project/trainers/PPOTrainer.py b/src/brittle_star_project/trainers/PPOTrainer.py index 81aa361..569f276 100644 --- a/src/brittle_star_project/trainers/PPOTrainer.py +++ b/src/brittle_star_project/trainers/PPOTrainer.py @@ -17,6 +17,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, @@ -132,7 +133,7 @@ def _linear_schedule(count, minibatch_count, update_epochs, num_iterations, lear def _normalize_obs(obs, mean, var, eps=1e-8): return jnp.clip((obs - mean) / jnp.sqrt(var + eps), -10.0, 10.0) - +# TODO: update to work with new obs_processor def _convert_obs_dict_to_array_morphology(obs_dict, morph_mode, num_segments: int, num_arms: int): @logged_jit def _filter_and_flatten(o) -> jnp.ndarray: @@ -203,7 +204,6 @@ _SEGMENT_SCALED_KEYS = frozenset( ) -# TODO: update to work with extra dimension + message passing def _get_action_and_value_noise( sensor: nn.Module, feature_extractor: nn.Module, @@ -329,7 +329,7 @@ def _reward_fn(env_state, next_env_state): def _step_env_wrapped( - episode_stats, env_state, action, env_step_fn, morph_mode, num_segments: int, num_arms: int + episode_stats, env_state, action, env_step_fn, morph_mode, num_segments: int, num_arms: int, # TODO: obs_processor ): next_env_state = env_step_fn(env_state, action) @@ -361,6 +361,7 @@ def _step_env_wrapped( reward, done, ), + # TODO (obs_processor(next_env_state.observations), reward, done), ) @@ -501,8 +502,8 @@ class PPOTrainer: self.key = jax.random.PRNGKey(self.experiment.seed) self.morph_mode = self.cfg.morphology.morph_mode + self.segments_per_arm = jnp.asarray(self.cfg.morphology.segments_per_arm, dtype=jnp.int32) - self.num_segments = self.segments_per_arm.sum().item() self.num_arms = jnp.where(self.segments_per_arm > 0, 1, 0).sum().item() @@ -523,6 +524,12 @@ class PPOTrainer: self.actor.apply = logged_jit(self.actor.apply) self.critic.apply = logged_jit(self.critic.apply) + # 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, + ) + action_low = jnp.asarray(self.env.single_action_space.low, dtype=jnp.float32) action_high = jnp.asarray(self.env.single_action_space.high, dtype=jnp.float32) @@ -536,6 +543,7 @@ class PPOTrainer: morph_mode=self.morph_mode, num_segments=self.num_segments, num_arms=self.num_arms, + # TODO: obs_processor=self.obs_processor, ), sensor=self.sensor, feature_extractor=self.feature_extractor, @@ -570,9 +578,6 @@ class PPOTrainer: def apply_feature(p, x): return apply_shared(self.feature_extractor, p, x) - def apply_message_passer(p, x): - return apply_shared(self.message_passer, p, x) - self._ppo = PPO(self.ppo, apply_sensor, apply_actor, apply_critic, apply_feature) self.agent_state = self._init_agent_state() @@ -623,6 +628,7 @@ class PPOTrainer: ) dummy_reset = self.env.reset(seed=0) + for k, v in dummy_reset.observations.items(): self.logger.debug(k, v.shape) sample_obs = _convert_obs_dict_to_array_morphology( @@ -697,6 +703,7 @@ class PPOTrainer: critic_params = self.critic.init(critic_key, critic_input) self.logger.debug( f"[_init_agent_state] critic_params: {jax.tree.map(lambda x: x.shape, critic_params)}" + # TODO: sample_obs = self.obs_processor(dummy_reset.observations)[0] # take first env ) return TrainState.create( @@ -924,6 +931,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_morphology( env_state.observations, self.morph_mode, @@ -931,6 +939,7 @@ class PPOTrainer: self.num_arms, ) self.logger.info(f"[train] next_obs: {next_obs.shape}") + # TODO: 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()}") @@ -945,7 +954,7 @@ class PPOTrainer: env_state, next_obs, next_done, training_measurements, storage = self._step( env_state, next_obs, next_done, iteration=iteration ) - self.logger.info(f"[train] next_obs (post-step): {next_obs.shape}") + self.logger.debug(f"[train] next_obs (post-step): {next_obs.shape}") self._update_obs_stats(next_obs) next_obs = _normalize_obs(next_obs, self.obs_mean, self.obs_var) diff --git a/tests/test_evaluation.py b/tests/test_evaluation.py new file mode 100644 index 0000000..bbb03f8 --- /dev/null +++ b/tests/test_evaluation.py @@ -0,0 +1,114 @@ +import numpy as np +import pytest +import yaml +from pathlib import Path + +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, + 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) + + +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) 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}") 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), ( 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}." + ) 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"