A. K. Mackay
Contents

A.2.2

Appendix A.2 Reinforcement learning

Dynamics ensembles for model-based RL

Neural-network ensembles that learn how an environment evolves from offline data, comparing deterministic and probabilistic models. They are the learned models behind J-MOReL and MBiCEM.

2025

Part of my MSc thesis

PyTorch, D4RL, MLflow, Guild AI, Hugging Face Hub

What it does

Model-based agents plan or train inside a learned model of the environment, so both the model’s accuracy and its sense of its own uncertainty matter. This library trains ensembles of neural networks on D4RL offline datasets and compares deterministic networks with probabilistic ones that predict a Gaussian over the next state.

Findings

On halfcheetah-medium the deterministic ensemble reaches the lower training loss, but the probabilistic ensemble generalises better and reaches the lowest validation loss. Without a learning-rate scheduler, the probabilistic ensemble’s training becomes unstable.

Engineering

Metrics and artifacts are tracked in MLflow and experiments are staged with Guild AI. Trained ensembles are published on Hugging Face and reused directly by J-MOReL and MBiCEM.

Diagram of three neural networks, each taking the current state and action of a HalfCheetah robot and predicting its next state.
Each network in the ensemble predicts the next state from the current state and action. Where the predictions disagree, the model is uncertain.
Line chart of validation MSE over 500 epochs for three ensembles. The probabilistic ensemble with a scheduler ends lowest, near 0.0023; without a scheduler it is noisy, around 0.0027; the deterministic ensemble plateaus near 0.0031.Line chart of validation MSE over 500 epochs for three ensembles. The probabilistic ensemble with a scheduler ends lowest, near 0.0023; without a scheduler it is noisy, around 0.0027; the deterministic ensemble plateaus near 0.0031.
Validation MSE on halfcheetah-medium. The probabilistic ensemble generalises best; without a learning-rate scheduler its loss becomes unstable, and the deterministic ensemble plateaus highest.