diff --git a/README.md b/README.md index 572fd73..e9cabc6 100644 --- a/README.md +++ b/README.md @@ -13,6 +13,21 @@ To set up the UV module, you can run the following command: uv sync --frozen ``` +## Repository Structure + +```text +. +├── configs/ # Hydra configuration files (YAML) +├── docs/ # Comprehensive documentation and API guides +├── runs/ # Default output directory for Hydra and training artifacts +├── scripts/ # High-level entrypoints for training, simulation, and evaluation +├── src/ +│ └── brittle_star_project/ # Core library and environment logic +│ ├── evaluation/ # Checkpoint evaluation, rollout logic, and metrics persistence +│ └── trainers/ # Training loop implementations (e.g., PPO) +└── tests/ # Unit and integration tests +``` + ## Usage For detailed instructions on how to use the project, please refer to the **[API Documentation](docs/README.md)**. diff --git a/docs/api/evaluation.md b/docs/api/evaluation.md new file mode 100644 index 0000000..a47d056 --- /dev/null +++ b/docs/api/evaluation.md @@ -0,0 +1,47 @@ +# Checkpoint & Model Evaluation + +This guide covers how to evaluate trained brittle star models, both during training and as a post-training analysis step. + +## Checkpoint Evaluation (During Training) + +The `PPOTrainer` can automatically evaluate every saved checkpoint using the fast MJX backend. This is enabled via configuration. + +### Configuration + +In your experiment config or via CLI: +```bash +python scripts/train.py evaluation.evaluate_checkpoints=true evaluation.eval_max_steps=5000 +``` + +Results are saved to `runs//metrics/checkpoint_evaluation.csv` and synced to Weights & Biases if enabled. + +## Cross-Model Comparison + +To compare different architectures or runs, use `scripts/compare_models.py`. + +1. Create or update a YAML file in `configs/evaluation/`. +2. Run the Comparison: + +```bash +python scripts/compare_models.py evaluation=poster +``` + +The script will run the specified number of episodes for each model and produce a single CSV with return and velocity metrics. + +## Post-hoc Checkpoint Evaluation + +If you didn't enable evaluation during training, or want to re-run it with different settings, use `scripts/evaluate_checkpoints.py`. + +```bash +python scripts/evaluate_checkpoints.py \ + simulation.model_path=runs//final_model.flax \ + evaluation.eval_seed=42 +``` + +This script scans the `checkpoints/` directory and evaluates every `.flax` file it finds. + +## Metrics Explained + +- **`eval_return`**: The accumulated shaped reward using the `reward_fn` defined in `PPOTrainer`. +- **`approx_max_velocity`**: Calculated as `(initial_dist - final_dist) / total_steps`. Note that this is an average velocity over the episode. +- **`reached_target`**: Boolean indicating if the robot reached the target within the max steps. diff --git a/docs/api/simulation.md b/docs/api/simulation.md index 0a0e041..1e288b9 100644 --- a/docs/api/simulation.md +++ b/docs/api/simulation.md @@ -37,3 +37,5 @@ uv run scripts/simulate.py \ Videos and evaluation metadata are stored in timestamped folders alongside the model: `runs/your_run/final_model_evaluations/eval_/simulation.mp4` + +For batch evaluation and cross-model comparison, see the **[Evaluation Guide](./evaluation.md)**. diff --git a/docs/api/training.md b/docs/api/training.md index a9678bb..d3e6c51 100644 --- a/docs/api/training.md +++ b/docs/api/training.md @@ -38,12 +38,18 @@ To run with your custom experiment file: 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 ``` +## Evaluation During Training + +By default, the trainer saves checkpoints but does not evaluate them. To enable automatic headless evaluation of every saved checkpoint, set `evaluation.evaluate_checkpoints=true`: + +```bash +uv run python scripts/train.py evaluation.evaluate_checkpoints=true +``` + +For more details on evaluation metrics and comparison tools, see [Evaluation](./evaluation.md). + For more details on tracking your experiments, see [Tracking & Monitoring](./tracking.md).