# Environment Transformer and Policy Optimization for Model-Based Offline Reinforcement Learning

arXiv: https://arxiv.org/abs/2303.03811 · HTML: https://mlsanity.com/paper/2303.03811/

## Key facts

- Cited by 4 papers in the map (data through 14 Oct 2026).
- Placed in Offline RL algorithms, one of 13 directions in Offline RL and imitation learning from data.
- Builds on 30 of its 40 references that are in the map.
- First posted on arXiv 7 Mar 2023.

## Placed in

- Offline RL and imitation learning from data / [Offline RL algorithms](https://mlsanity.com/offlinerl/offline-rl-algorithms/)

## Builds on (most cited first, 30 of 30)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Soft Actor-Critic: Off-Policy Maximum Entropy Deep Reinforcement Learning with a Stochastic Actor](https://mlsanity.com/paper/1801.01290/) (2018)
- [Sequence to Sequence Learning with Neural Networks](https://mlsanity.com/paper/1409.3215/) (2014)
- [High-Dimensional Continuous Control Using Generalized Advantage Estimation](https://mlsanity.com/paper/1506.02438/) (2015)
- [Conservative Q-Learning for Offline Reinforcement Learning](https://mlsanity.com/paper/2006.04779/) (2020)
- [Decision Transformer: Reinforcement Learning via Sequence Modeling](https://mlsanity.com/paper/2106.01345/) (2021)
- [D4RL: Datasets for Deep Data-Driven Reinforcement Learning](https://mlsanity.com/paper/2004.07219/) (2020)
- [Offline Reinforcement Learning with Implicit Q-Learning](https://mlsanity.com/paper/2110.06169/) (2021)
- [Off-Policy Deep Reinforcement Learning without Exploration](https://mlsanity.com/paper/1812.02900/) (2018)
- [A Minimalist Approach to Offline Reinforcement Learning](https://mlsanity.com/paper/2106.06860/) (2021)
- [Deep Reinforcement Learning in a Handful of Trials using Probabilistic Dynamics Models](https://mlsanity.com/paper/1805.12114/) (2018)
- [When to Trust Your Model: Model-Based Policy Optimization](https://mlsanity.com/paper/1906.08253/) (2019)
- [AWAC: Accelerating Online Reinforcement Learning with Offline Datasets](https://mlsanity.com/paper/2006.09359/) (2020)
- [Behavior Regularized Offline Reinforcement Learning](https://mlsanity.com/paper/1911.11361/) (2019)
- [Offline Reinforcement Learning as One Big Sequence Modeling Problem](https://mlsanity.com/paper/2106.02039/) (2021)
- [MOPO: Model-based Offline Policy Optimization](https://mlsanity.com/paper/2005.13239/) (2020)
- [MOReL : Model-Based Offline Reinforcement Learning](https://mlsanity.com/paper/2005.05951/) (2020)
- [COMBO: Conservative Offline Model-Based Policy Optimization](https://mlsanity.com/paper/2102.08363/) (2021)
- [Critic Regularized Regression](https://mlsanity.com/paper/2006.15134/) (2020)
- [PLAS: Latent Action Space for Offline Reinforcement Learning](https://mlsanity.com/paper/2011.07213/) (2020)
- [RAMBO-RL: Robust Adversarial Model-Based Offline Reinforcement Learning](https://mlsanity.com/paper/2204.12581/) (2022)
- [Colossal-AI: A Unified Deep Learning System For Large-Scale Parallel Training](https://mlsanity.com/paper/2110.14883/) (2021)
- [Hierarchical Planning Through Goal-Conditioned Offline Reinforcement Learning](https://mlsanity.com/paper/2205.11790/) (2022)
- [Offline Reinforcement Learning with Reverse Model-based Imagination](https://mlsanity.com/paper/2110.00188/) (2021)
- [Offline Meta-Reinforcement Learning for Industrial Insertion](https://mlsanity.com/paper/2110.04276/) (2021)
- [Offline Reinforcement Learning with Soft Behavior Regularization](https://mlsanity.com/paper/2110.07395/) (2021)
- [RACER: Rapid Collaborative Exploration with a Decentralized Multi-UAV System](https://mlsanity.com/paper/2209.08533/) (2022)
- [Teach-Repeat-Replan: A Complete and Robust System for Aggressive Flight in Complex Environments](https://mlsanity.com/paper/1907.00520/) (2019)
- [How to Spend Your Robot Time: Bridging Kickstarting and Offline Reinforcement Learning for Vision-based Robotic Manipulation](https://mlsanity.com/paper/2205.03353/) (2022)
- [Offline Learning of Counterfactual Predictions for Real-World Robotic Reinforcement Learning](https://mlsanity.com/paper/2011.05857/) (2020)

## Built on by (most cited first, 4 of 4)

- [Deep Generative Models for Offline Policy Learning: Tutorial, Survey, and Perspectives on Future Directions](https://mlsanity.com/paper/2402.13777/) (2024)
- [Deep Reinforcement Learning for Time-Critical Wilderness Search And Rescue Using Drones](https://mlsanity.com/paper/2405.12800/) (2024)
- [Model-Based Offline Reinforcement Learning with Adversarial Data Augmentation](https://mlsanity.com/paper/2503.20285/) (2025)
- [Guided Riemannian Optimization (GuRO): Bridging Model Predictive Control and Decision Transformers](https://mlsanity.com/paper/2608.23204/) (2026)

## Abstract (arXiv)

Interacting with the actual environment to acquire data is often costly and time-consuming in robotic tasks. Model-based offline reinforcement learning (RL) provides a feasible solution. On the one hand, it eliminates the requirements of interaction with the actual environment. On the other hand, it learns the transition dynamics and reward function from the offline datasets and generates simulated rollouts to accelerate training. Previous model-based offline RL methods adopt probabilistic ensemble neural networks (NN) to model aleatoric uncertainty and epistemic uncertainty. However, this results in an exponential increase in training time and computing resource requirements. Furthermore, these methods are easily disturbed by the accumulative errors of the environment dynamics models when simulating long-term rollouts. To solve the above problems, we propose an uncertainty-aware sequence modeling architecture called Environment Transformer. It models the probability distribution of the environment dynamics and reward function to capture aleatoric uncertainty and treats epistemic uncertainty as a learnable noise parameter. Benefiting from the accurate modeling of the transition dynamics and reward function, Environment Transformer can be combined with arbitrary planning, dynamics programming, or policy optimization algorithms for offline RL. In this case, we perform Conservative Q-Learning (CQL) to learn a conservative Q-function. Through simulation experiments, we demonstrate that our method achieves or exceeds state-of-the-art performance in widely studied offline RL benchmarks. Moreover, we show that Environment Transformer's simulated rollout quality, sample efficiency, and long-term rollout simulation capability are superior to those of previous model-based offline RL methods.