Multi-Step Model-Based Reinforcement Learning

A collection of agents and models for multi-step model-based reinforcement learning (MBRL) algorithms.

Author: Abdelhakim Benechehab

Overview

This project implements a framework for model-based reinforcement learning with support for:

• Multiple model architectures (neural networks, mixture density networks, XGBoost, etc.)
• Multi-step prediction horizons with various loss functions
• Different planning strategies (Random Shooting, Soft Actor-Critic)
• Hyperparameter optimization using HEBO
• Integration with dm_control environments (CartPole Swingup, HalfCheetah)

Project Structure

multi_step_mbrl/
├── src/multi_step_mbrl/
│   ├── exp_launcher.py          # Main experiment launcher (entry point)
│   ├── config/                  # Hydra configuration files
│   │   ├── main.yaml            # Main configuration
│   │   ├── mode/                # Execution modes (dynamic, static, hyperopt)
│   │   ├── model/               # Model configurations
│   │   ├── actor/               # Actor configurations
│   │   └── environment/         # Environment configurations
│   ├── models/                  # Model implementations
│   │   ├── dn/                  # Deterministic Networks
│   │   ├── mdn/                 # Mixture Density Networks
│   │   ├── nn/                  # Standard Neural Networks
│   │   ├── multi_step_nn/       # Multi-step Neural Networks
│   │   ├── xgboost/             # XGBoost model
│   │   └── ...                  # Other model variants
│   ├── actors/                  # Actor/planner implementations
│   │   ├── RS/                  # Random Shooting
│   │   └── sac/                 # Soft Actor-Critic
│   ├── envs/                    # Environment definitions
│   │   ├── cartpole_swingup/    # CartPole Swingup task
│   │   └── halfcheetah_with_reward/
│   ├── trainers/                # Training utilities
│   │   ├── loss_fn.py           # Loss functions (MSE, NLL, MultiStep)
│   │   ├── data_processing.py   # Multi-step data preprocessing
│   │   └── pytorch_utils.py     # PyTorch utilities
│   └── visualization/           # Analysis and visualization tools
├── pyproject.toml               # Project metadata
└── requirements.txt             # Dependencies

Installation

Requirements

• Python >= 3.8
• PyTorch (CPU)
• MuJoCo 2.3.3
• dm_control 1.0.11

Install

# Clone the repository
git clone <repository-url>
cd multi_step_mbrl

# Install dependencies
pip install -e .

# For development
pip install -e ".[dev]"

Main Script: exp_launcher.py

The main entry point is src/multi_step_mbrl/exp_launcher.py. It uses Hydra for configuration management and supports three execution modes.

Usage

cd src/multi_step_mbrl
python exp_launcher.py trial_name=<your_trial_name> [options]

Core Functions

Function	Description
launch()	Hydra-decorated main entry point that orchestrates experiment execution
mbrl_generate_exp()	Generates experiment directory structure with model/actor hyperparameters
mbrl_run_command()	Executes the MBRL training loop with environment interaction
simulate_command()	Runs long-horizon prediction simulations for evaluation

Execution Modes

1. Dynamic Mode (default)

Online MBRL training with agent-environment interaction.

python exp_launcher.py trial_name=my_trial mode=dynamic

Key parameters:

• n_epochs: Number of training epochs
• episodic_update: Update model after each episode
• num_envs: Number of parallel environments
• use_tensorboard: Enable TensorBoard logging
• save_model: Save trained model checkpoints

2. Static Mode

Offline model evaluation with cross-validation.

python exp_launcher.py trial_name=my_trial mode=static mode.data_label=train

Key parameters:

• data_label: Dataset label for evaluation
• simulate_traces: Run long-horizon trace simulation
• future_length: Prediction horizon length for simulation

3. Hyperopt Mode

Automated hyperparameter optimization using HEBO.

python exp_launcher.py trial_name=my_trial mode=hyperopt mode.n_trials=50

Key parameters:

• engine_name: Optimization engine (default: "random")
• n_trials: Number of optimization trials
• max_concurrent_runs: Parallel trial execution

Configuration Options

Model Selection

python exp_launcher.py trial_name=my_trial model=<model_name>

Available models: dn, mdn, nn, multi_step_nn, linear, xgboost, baseline, etc.

Actor Selection

python exp_launcher.py trial_name=my_trial actor=<actor_name>

Available actors: RS (Random Shooting), sac (Soft Actor-Critic)

Environment Selection

python exp_launcher.py trial_name=my_trial environment=<env_name>

Available environments: cartpole_swingup, halfcheetah_with_reward

Example Commands

# Basic dynamic training with DN model
python exp_launcher.py trial_name=exp001 model=dn actor=RS environment=cartpole_swingup

# Static evaluation with long-horizon simulation
python exp_launcher.py trial_name=exp001 mode=static mode.data_label=train mode.future_length=100

# Hyperparameter optimization
python exp_launcher.py trial_name=exp001 mode=hyperopt mode.n_trials=100 mode.max_concurrent_runs=4

# Debug mode (reduced epochs and timesteps)
python exp_launcher.py trial_name=debug_run debug=True

# Custom hyperparameters
python exp_launcher.py trial_name=exp001 model.hypers.n_epochs=200 actor.hypers.planning_horizon=15

# Set random seed
python exp_launcher.py trial_name=exp001 seed=42

Output Structure

Experiments are saved in the following structure:

experiments/<environment>/trials/<trial_name>/
├── model/                    # Model files and hyperparameters
│   ├── generative_regressor.py
│   ├── feature_extractor.py
│   └── training_output/      # Trained models and metrics
├── actor/                    # Actor files and hyperparameters
│   └── actor.py
└── hydra/                    # Hydra configuration logs

Models

Model	Description
dn	Deterministic Networks
mdn	Mixture Density Networks (probabilistic)
nn	Standard Neural Network
multi_step_nn	Multi-step prediction Neural Network
dn_multi_step_loss_beta	DN with beta-weighted multi-step loss
dn_multi_step_loss_free	DN with free multi-step loss weighting
linear	Linear regressor
xgboost	XGBoost gradient boosting
baseline	Baseline model

Dependencies

Core:

• numpy, pandas, torch
• hydra-core (configuration)
• stable-baselines3 (RL algorithms)
• dm_control, gymnasium, mujoco (environments)

Internal packages:

• ramp-workflow - ML workflow framework
• ramp-hyperopt - Hyperparameter optimization
• mbrltools - MBRL utilities

Development

# Install dev dependencies
pip install -e ".[dev]"

# Run tests
pytest

# Run linter
ruff check src/

# Format code
ruff format src/

License

See LICENSE file for details.