Merge pull request #45 from SELab-3-2026/fix/input-space
This commit is contained in:
GitHub
commit
70ef9d4058
37 changed files with 1665 additions and 934 deletions
44
README.md
44
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
|
||||
|
||||
|
|
|
|||
|
|
@ -11,6 +11,7 @@ defaults:
|
|||
- morphology: 5_arms_full
|
||||
- arena: default
|
||||
- environment: directed_locomotion
|
||||
- obs_bounds: default
|
||||
- simulation: default
|
||||
- _self_
|
||||
|
||||
|
|
|
|||
1
configs/obs_bounds/default.yaml
Normal file
1
configs/obs_bounds/default.yaml
Normal file
|
|
@ -0,0 +1 @@
|
|||
# Defaults provided by dataclass
|
||||
|
|
@ -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 `<model_name>_metadata.yaml` alongside the model_path.
|
||||
metadata_path: null
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
|
|
|||
|
|
@ -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
|
||||
```
|
||||
39
docs/api/simulation.md
Normal file
39
docs/api/simulation.md
Normal file
|
|
@ -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_<timestamp>/simulation.mp4`
|
||||
60
docs/api/tracking.md
Normal file
60
docs/api/tracking.md
Normal file
|
|
@ -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).
|
||||
49
docs/api/training.md
Normal file
49
docs/api/training.md
Normal file
|
|
@ -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).
|
||||
|
|
@ -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
|
||||
```
|
||||
|
|
|
|||
|
|
@ -38,6 +38,10 @@ cuda = [
|
|||
analysis = [
|
||||
"tensorboard",
|
||||
]
|
||||
evaluation = [
|
||||
"imageio>=2.35.0",
|
||||
"imageio-ffmpeg>=0.5.1",
|
||||
]
|
||||
|
||||
[dependency-groups]
|
||||
dev = [
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
141
scripts/tools/dump_mjcf.py
Normal file
141
scripts/tools/dump_mjcf.py
Normal file
|
|
@ -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/<name>.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()
|
||||
138
scripts/tools/extract_observation_bounds.py
Normal file
138
scripts/tools/extract_observation_bounds.py
Normal file
|
|
@ -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()
|
||||
|
|
@ -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",
|
||||
]
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -1,7 +1,9 @@
|
|||
from .env_config import ArenaConfig, EnvConfig, MorphologyConfig
|
||||
from .env_types import Backend, Task
|
||||
from .env_wrapper import BrittleStarEnv, StepResult
|
||||
from .env_wrapper import BrittleStarEnv
|
||||
from .factory import BrittleStarEnvFactory
|
||||
from .obs_processing import create_obs_processor
|
||||
from .padded_obs_wrapper import compute_padding_masks
|
||||
|
||||
__all__ = [
|
||||
"ArenaConfig",
|
||||
|
|
@ -10,6 +12,7 @@ __all__ = [
|
|||
"Backend",
|
||||
"Task",
|
||||
"BrittleStarEnv",
|
||||
"StepResult",
|
||||
"BrittleStarEnvFactory",
|
||||
"create_obs_processor",
|
||||
"compute_padding_masks",
|
||||
]
|
||||
|
|
|
|||
|
|
@ -60,3 +60,31 @@ class EnvConfig:
|
|||
# Light escape
|
||||
# Per docs in upstream env config: integer factors of 200.
|
||||
light_perlin_noise_scale: int = 0
|
||||
|
||||
|
||||
@dataclass
|
||||
class ObservationBoundsConfig:
|
||||
"""Physical observation bounds for deterministic min-max normalization."""
|
||||
|
||||
# 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),
|
||||
}
|
||||
|
|
|
|||
91
src/brittle_star_project/environment/obs_processing.py
Normal file
91
src/brittle_star_project/environment/obs_processing.py
Normal file
|
|
@ -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))
|
||||
|
|
@ -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
|
||||
|
|
|
|||
21
src/brittle_star_project/evaluation/__init__.py
Normal file
21
src/brittle_star_project/evaluation/__init__.py
Normal file
|
|
@ -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",
|
||||
]
|
||||
107
src/brittle_star_project/evaluation/checkpoint.py
Normal file
107
src/brittle_star_project/evaluation/checkpoint.py
Normal file
|
|
@ -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,
|
||||
)
|
||||
89
src/brittle_star_project/evaluation/policy.py
Normal file
89
src/brittle_star_project/evaluation/policy.py
Normal file
|
|
@ -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()
|
||||
163
src/brittle_star_project/evaluation/rollout.py
Normal file
163
src/brittle_star_project/evaluation/rollout.py
Normal file
|
|
@ -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}"
|
||||
)
|
||||
148
src/brittle_star_project/evaluation/video.py
Normal file
148
src/brittle_star_project/evaluation/video.py
Normal file
|
|
@ -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,
|
||||
)
|
||||
|
|
@ -1,3 +0,0 @@
|
|||
from .renderer import simulate_policy, SimulationConfig, ControlPolicy
|
||||
|
||||
__all__ = ["simulate_policy", "SimulationConfig", "ControlPolicy"]
|
||||
|
|
@ -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)
|
||||
|
|
@ -16,6 +16,7 @@ from experiment_logger import get_logger
|
|||
from brittle_star_project.configs.main_config import BrittleStarConfig
|
||||
from brittle_star_project.dataclasses import EpisodeStatistics
|
||||
from brittle_star_project.environment.BrittleStarJaxEnvWrapper import BrittleStarJaxEnvWrapper
|
||||
from brittle_star_project.environment.obs_processing import create_obs_processor
|
||||
from brittle_star_project.MLPs.mlps import (
|
||||
Actor,
|
||||
AgentParams,
|
||||
|
|
@ -25,19 +26,6 @@ from brittle_star_project.MLPs.mlps import (
|
|||
)
|
||||
from brittle_star_project.ppo import PPO
|
||||
|
||||
# TODO: move to config
|
||||
_ALLOWED_OBS_KEYS = {
|
||||
"joint_position",
|
||||
"joint_velocity",
|
||||
"joint_actuator_force",
|
||||
"actuator_force",
|
||||
"disk_position",
|
||||
"disk_rotation",
|
||||
"disk_linear_velocity",
|
||||
"disk_angular_velocity",
|
||||
"unit_xy_direction_to_target",
|
||||
"xy_distance_to_target",
|
||||
}
|
||||
# TODO: clip scaled reward?
|
||||
|
||||
|
||||
|
|
@ -65,27 +53,6 @@ def _linear_schedule(count, minibatch_count, update_epochs, num_iterations, lear
|
|||
return learning_rate * frac
|
||||
|
||||
|
||||
@jax.jit
|
||||
def _normalize_obs(obs, mean, var, eps=1e-8):
|
||||
return jnp.clip((obs - mean) / jnp.sqrt(var + eps), -10.0, 10.0)
|
||||
|
||||
|
||||
@jax.jit
|
||||
def _convert_obs_dict_to_array(obs_dict: dict) -> jnp.ndarray:
|
||||
"""Convert the raw observation dict → flat array, filtering unwanted keys."""
|
||||
|
||||
def _filter_and_flatten(o: dict) -> jnp.ndarray:
|
||||
values = []
|
||||
for key in sorted(o.keys()):
|
||||
if key in _ALLOWED_OBS_KEYS: # TODO: NORMALIZATION or .. of observations??
|
||||
v = o[key]
|
||||
if v.size > 0:
|
||||
values.append(jnp.asarray(v).flatten())
|
||||
return jnp.concatenate(values)
|
||||
|
||||
return jax.vmap(_filter_and_flatten)(obs_dict)
|
||||
|
||||
|
||||
def _get_action_and_value_noise(
|
||||
sensor: GenericDenseLayersWithActivation,
|
||||
feature_extractor: GenericDenseLayersWithActivation,
|
||||
|
|
@ -168,7 +135,7 @@ def _reward_fn(env_state, next_env_state):
|
|||
return jnp.where(next_env_state.terminated, 50.0, clipped_env_reward - penalty)
|
||||
|
||||
|
||||
def _step_env_wrapped(episode_stats, env_state, action, env_step_fn):
|
||||
def _step_env_wrapped(episode_stats, env_state, action, env_step_fn, obs_processor):
|
||||
next_env_state = env_step_fn(env_state, action)
|
||||
|
||||
reward = _reward_fn(env_state, next_env_state)
|
||||
|
|
@ -192,7 +159,7 @@ def _step_env_wrapped(episode_stats, env_state, action, env_step_fn):
|
|||
return (
|
||||
episode_stats,
|
||||
next_env_state,
|
||||
(_convert_obs_dict_to_array(next_env_state.observations), reward, done),
|
||||
(obs_processor(next_env_state.observations), reward, done),
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -303,6 +270,12 @@ class PPOTrainer:
|
|||
|
||||
self.key = jax.random.PRNGKey(self.experiment.seed)
|
||||
|
||||
# Build the centralized observation processor: derive -> normalize -> pad -> flatten.
|
||||
self.obs_processor = create_obs_processor(
|
||||
bounds_dict=self.cfg.obs_bounds.to_bounds_dict(),
|
||||
padding_masks=self.env.padding_masks,
|
||||
)
|
||||
|
||||
self.sensor, self.feature_extractor, self.actor, self.critic = self._init_agent()
|
||||
self.sensor.apply = jax.jit(self.sensor.apply)
|
||||
self.feature_extractor.apply = jax.jit(self.feature_extractor.apply)
|
||||
|
|
@ -316,7 +289,11 @@ class PPOTrainer:
|
|||
partial(
|
||||
_rollout_jit,
|
||||
max_steps=self.ppo.num_steps,
|
||||
step_env_fn=partial(_step_env_wrapped, env_step_fn=self.env.step),
|
||||
step_env_fn=partial(
|
||||
_step_env_wrapped,
|
||||
env_step_fn=self.env.step,
|
||||
obs_processor=self.obs_processor,
|
||||
),
|
||||
sensor=self.sensor,
|
||||
feature_extractor=self.feature_extractor,
|
||||
actor=self.actor,
|
||||
|
|
@ -367,10 +344,7 @@ class PPOTrainer:
|
|||
)
|
||||
|
||||
dummy_reset = self.env.reset(seed=0)
|
||||
sample_obs = _convert_obs_dict_to_array(dummy_reset.observations)[0] # take first env
|
||||
self.obs_mean = jnp.zeros((len(sample_obs),))
|
||||
self.obs_var = jnp.ones((len(sample_obs),))
|
||||
self.obs_count = 1e-4
|
||||
sample_obs = self.obs_processor(dummy_reset.observations)[0] # take first env
|
||||
sensor_params = self.sensor.init(sensor_key, sample_obs)
|
||||
feature_extractor_params = self.feature_extractor.init(feature_extractor_key, sample_obs)
|
||||
actor_params = self.actor.init(actor_key, self.sensor.apply(sensor_params, sample_obs))
|
||||
|
|
@ -410,25 +384,6 @@ class PPOTrainer:
|
|||
returned_episode_lengths=jnp.zeros(self.ppo.num_envs, dtype=jnp.int32),
|
||||
)
|
||||
|
||||
def _update_obs_stats(self, obs: jnp.ndarray):
|
||||
batch_mean = jnp.mean(obs, axis=0)
|
||||
batch_var = jnp.var(obs, axis=0)
|
||||
batch_count = obs.shape[0]
|
||||
|
||||
delta = batch_mean - self.obs_mean
|
||||
total_count = self.obs_count + batch_count
|
||||
|
||||
new_mean = self.obs_mean + delta * batch_count / total_count
|
||||
|
||||
m_a = self.obs_var * self.obs_count
|
||||
m_b = batch_var * batch_count
|
||||
M2 = m_a + m_b + delta**2 * self.obs_count * batch_count / total_count
|
||||
new_var = M2 / total_count
|
||||
|
||||
self.obs_mean = new_mean
|
||||
self.obs_var = new_var
|
||||
self.obs_count = total_count
|
||||
|
||||
def _rollout(self, env_state, next_obs, next_done) -> tuple[Any, ...]:
|
||||
return self._rollout_jit(
|
||||
self.agent_state,
|
||||
|
|
@ -594,7 +549,7 @@ class PPOTrainer:
|
|||
self.logger.log_non_interactive(f"Initial reset started: {time.ctime()}")
|
||||
|
||||
env_state = self.env.reset(seed=self.experiment.seed)
|
||||
next_obs = _convert_obs_dict_to_array(env_state.observations)
|
||||
next_obs = self.obs_processor(env_state.observations)
|
||||
next_done = jnp.zeros(self.ppo.num_envs, dtype=jnp.bool_)
|
||||
|
||||
self.logger.log_non_interactive(f"Initial reset completed: {time.ctime()}")
|
||||
|
|
@ -609,8 +564,6 @@ class PPOTrainer:
|
|||
env_state, next_obs, next_done, training_measurements, storage = self._step(
|
||||
env_state, next_obs, next_done, iteration=iteration
|
||||
)
|
||||
self._update_obs_stats(next_obs)
|
||||
next_obs = _normalize_obs(next_obs, self.obs_mean, self.obs_var)
|
||||
|
||||
global_step += self.ppo.num_steps * self.ppo.num_envs
|
||||
self._log(
|
||||
|
|
|
|||
114
tests/test_evaluation.py
Normal file
114
tests/test_evaluation.py
Normal file
|
|
@ -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)
|
||||
|
|
@ -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}")
|
||||
|
||||
|
|
|
|||
|
|
@ -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), (
|
||||
|
|
|
|||
78
tests/test_target_direction.py
Normal file
78
tests/test_target_direction.py
Normal file
|
|
@ -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}."
|
||||
)
|
||||
22
uv.lock
generated
22
uv.lock
generated
|
|
@ -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"
|
||||
|
|
|
|||
Reference in a new issue