This is the pre-proceedings for the RLC 2026. You may expect minor changes.

Stable Planning through Aligned Representations in Model-Based Reinforcement Learning

By Misagh Soltani, and Forest Agostinelli

Reinforcement Learning Journal, vol. 7, 2026, pp. TBD.

Will be presented at the Reinforcement Learning Conference (RLC), MontrĂ©al, Quebec, Canada, August 15–17, 2026.


Download:

Keywords: Model-Based Reinforcement Learning, Planning, Reinforcement Learning, Alignment

Abstract:

World models that reliably predict state transitions can facilitate the integration of planning with reinforcement learning to solve sparse-reward, long-horizon sequential decision-making problems. However, enabling such learned world models to generalize to noise for which the underlying Markov decision process is invariant (e.g., lighting variations, blurring, or rotations) requires training the world model on a dataset with such noise. Unfortunately, this does not explicitly capture these invariances and requires that the world model, and any functions trained with data generated from it, be re-trained if it fails to generalize to new kinds of noise. To solve this problem, we introduce Stable Planning through Aligned Representations (SPAR), an algorithm that explicitly captures invariances by training an alignment network to map noisy states to the latent representation of clean states while keeping the learned world model fixed. SPAR first trains a world model with a discrete latent space on clean states. SPAR then trains an alignment network to map noisy states to the same latent representation as the clean states, exploiting the discrete latent structure by rounding the output of the alignment network after training so that it exactly matches the clean latent representation. As a result, SPAR preserves the ability to use a discrete world model across thousands of timesteps without accumulating any error by rounding the predicted discrete latent states. Our experiments show that SPAR can be used to successfully solve over 89\% and 96\% of problem instances from the Rubik's Cube and Sokoban domains across 16 and 18 different visual transformations respectively. Furthermore, SPAR can be used to solve the Rubik's Cube based on real-world pictures taken by a phone, without ever having seen real-world Rubik's Cubes during training. Our experiments show that continuous representations lead to model degradation, indicating that the discrete latent component is crucial to SPAR's success.


Citation Information:

Misagh Soltani and Forest Agostinelli. "Stable Planning through Aligned Representations in Model-Based Reinforcement Learning." Reinforcement Learning Journal, vol. 7, 2026, pp. TBD.

BibTeX:
@article{soltani2026stable,
    title={Stable Planning through Aligned Representations in Model-Based Reinforcement Learning},
    author={Misagh Soltani and Forest Agostinelli},
    journal={Reinforcement Learning Journal},
    volume={7},
    pages={},
    year={2026}
}