Predictable MDP Abstraction for Unsupervised Model-Based RL

TL;DR

This paper introduces Predictable MDP Abstraction (PMA), an unsupervised method that transforms complex MDPs into simpler, predictable forms to improve model-based RL performance without additional environment interactions.

Contribution

The paper proposes a novel unsupervised approach to learn an action space transformation that simplifies MDP prediction tasks, enabling more accurate models and zero-shot downstream control.

Findings

PMA significantly improves model accuracy over prior methods.

PMA enables zero-shot control on benchmark tasks.

The approach is theoretically sound and empirically validated.

Abstract

A key component of model-based reinforcement learning (RL) is a dynamics model that predicts the outcomes of actions. Errors in this predictive model can degrade the performance of model-based controllers, and complex Markov decision processes (MDPs) can present exceptionally difficult prediction problems. To mitigate this issue, we propose predictable MDP abstraction (PMA): instead of training a predictive model on the original MDP, we train a model on a transformed MDP with a learned action space that only permits predictable, easy-to-model actions, while covering the original state-action space as much as possible. As a result, model learning becomes easier and more accurate, which allows robust, stable model-based planning or model-based RL. This transformation is learned in an unsupervised manner, before any task is specified by the user. Downstream tasks can then be solved with model-based control in a zero-shot fashion, without additional environment interactions. We theoretically analyze PMA and empirically demonstrate that PMA leads to significant improvements over prior unsupervised model-based RL approaches in a range of benchmark environments. Our code and videos are available at https://seohong.me/projects/pma/