Discovering Minimal Reinforcement Learning Environments

TL;DR

The paper tackles the inefficiency of RL training when agents are repeatedly evaluated in the same environment by introducing synthetic, neural-network–parameterized environments (SEs) that are optimized to maximize transfer to a fixed EE. It extends prior work with an algorithm- and hyperparameter-invariant meta-learning approach, a horizon-curriculum to scale to Brax-like tasks, hardware-accelerated vectorized RL, and a surprising finding that contextual bandits can serve as effective proxies for complex MDPs while offering interpretability. Key contributions include showing CBs arise naturally from synthetic MDPs, enabling generalization across unseen algorithms and hyperparameters, and demonstrating downstream speedups in Learned Policy Optimization and other meta-learning contexts. The approach yields fast, transferable RL pretraining signals, provides interpretable insights into evaluation environments, and offers a practical pathway to accelerate downstream RL research and applications, including pretraining, neural architecture search, and meta-learning.

Abstract

Reinforcement learning (RL) agents are commonly trained and evaluated in the same environment. In contrast, humans often train in a specialized environment before being evaluated, such as studying a book before taking an exam. The potential of such specialized training environments is still vastly underexplored, despite their capacity to dramatically speed up training. The framework of synthetic environments takes a first step in this direction by meta-learning neural network-based Markov decision processes (MDPs). The initial approach was limited to toy problems and produced environments that did not transfer to unseen RL algorithms. We extend this approach in three ways: Firstly, we modify the meta-learning algorithm to discover environments invariant towards hyperparameter configurations and learning algorithms. Secondly, by leveraging hardware parallelism and introducing a curriculum on an agent's evaluation episode horizon, we can achieve competitive results on several challenging continuous control problems. Thirdly, we surprisingly find that contextual bandits enable training RL agents that transfer well to their evaluation environment, even if it is a complex MDP. Hence, we set up our experiments to train synthetic contextual bandits, which perform on par with synthetic MDPs, yield additional insights into the evaluation environment, and can speed up downstream applications.

Figures (16)

• Figure 1: (1) Training process for synthetic contextual bandits. Firstly, an agent is trained in an SCB (blue), observing only an initial state and reward in each episode. After training, the agent is frozen and transferred to an evaluation environment. The achieved episodic return is the training signal to update the SCB (red). (2) Aggregated training results for challenging Brax environments. Training in the SCB yields policies that are competitive with EE experts, sometimes even outperforming them. The x-axis shows expert normalized performance as $(R - R_\text{random}) / (R_\text{expert} - R_\text{random})$, where $R_\text{random}$ and $R_\text{expert}$ are the episodic returns achieved in the EE by an expert and a random policy, respectively. For full training results, see appendix \ref{['tab:brax_retsults_table']}. (3) Training curves for EE versus SCBs, saving orders of magnitude of environment steps. A complete visualization is given in figure \ref{['fig:training_one_row']}. (4) Visualization of observation feature importances for Pendulum-v1. For details refer to section \ref{['sec:results-interpret']}.
• Figure 2: CBs are a discovered property of synthetic environments. Left: Fraction of episode lengths in a synthetic MDP. In most environments, more than 80% of episodes are terminated after a single time step. Episodes were generated using 50 episodes of agents trained with each algorithm. Right: Normalized performance of synthetic MDPs with decreasing maximum episode length. "max" refers to the maximum episode of the evaluation environment. Shown are the IQM scores and 95% confidence intervals of 20 training runs, aggregated over all algorithms (see section \ref{['sec:methods']}). Performance is normalized as $(R - R_\text{SCB}) / (R - R_\text{random})$ for each algorithm individually, where $R_\text{SCB}$ is the return in the EE after training in the SCB, and $R_\text{random}$ is the return of a random policy.
• Figure 3: Meta-learned SCBs generalize across hyperparameters and towards out-of-distribution RL algorithms. Each column shows the return of a policy after training in either the SCB (blue) or the EE directly (red), using sampled hyperparameters and 10,000 environment steps. Additionally, we ablate the meta-training algorithm by using fixed hyperparameters in the inner loop. When sampling hyperparameters during the evaluation, the ablated SCBs (green) perform worse. Additionally, the SCBs generalize towards agents not used in the inner loop (right of dashed line). For more details, refer to appenix \ref{['sec:generality_appendix']}.
• Figure 4: Curriculum and training speed for Brax environments. Left: Visualization of different curricula. Right: Influence of curricula on the meta-training progress for the hopper environment. For no curriculum, we show the IQM and 95% CIs over 9 seeds. One seed is shown per curriculum.
• Figure 5: Training curves for different algorithms on Brax environments. Training in an SCB is roughly two orders of magnitude faster. We show the IQM performance with 95% bootstrapped confidence intervals. Refer to appendix \ref{['sec:training_speed_comparison_brax']} for complete training curves.

Theorems & Definitions (3)

• Theorem 1
• Theorem 1
• proof