Source-linked AI summary

RAMBO-RL: Robust Adversarial Model-Based Offline Reinforcement Learning

Marc Rigter, Bruno Lacerda, Nick Hawes

arXiv:2204.12581v3cs.LGcs.AI

TL;DR

Offline RL must learn effective policies from fixed data while avoiding unreliable behavior outside the dataset. RAMBO solves this through adversarial model-based policy optimisation, and it reports state-of-the-art benchmark performance with theoretically grounded pessimism. Its strongest results are on noisy or suboptimal datasets, while performance is weaker on high-quality Medium-Expert and AntMaze tasks.

  • Problem

    Offline RL seeks effective policies from fixed datasets, but distributional shift makes learned models inaccurate and enables model exploitation outside covered state-action pairs.

  • Method

    RAMBO formulates offline RL as a two-player zero-sum game, alternately optimizing the policy and an adversarial transition model to enforce conservatism.

  • Results

    RAMBO outperforms current state-of-the-art algorithms on standard D4RL benchmarks, with strongest overall scores on MuJoCo locomotion, Medium, and Medium-Replay datasets.

  • Takeaways & Limitations

    Adversarial model training provides a practical, theoretically grounded approach to conservative model-based offline RL.

  • Takeaways & Limitations

    RAMBO is less suitable for high-quality Medium-Expert datasets and, like other model-based methods, performs poorly on AntMaze domains.

Abstract

from arXiv · show

Offline reinforcement learning (RL) aims to find performant policies from logged data without further environment interaction. Model-based algorithms, which learn a model of the environment from the dataset and perform conservative policy optimisation within that model, have emerged as a promising approach to this problem. In this work, we present Robust Adversarial Model-Based Offline RL (RAMBO), a novel approach to model-based offline RL. We formulate the problem as a two-player zero sum game against an adversarial environment model. The model is trained to minimise the value function while still accurately predicting the transitions in the dataset, forcing the policy to act conservatively in areas not covered by the dataset. To approximately solve the two-player game, we alternate between optimising the policy and adversarially optimising the model. The problem formulation that we address is theoretically grounded, resulting in a probably approximately correct (PAC) performance guarantee and a pessimistic value function which lower bounds the value function in the true environment. We evaluate our approach on widely studied offline RL benchmarks, and demonstrate that it outperforms existing state-of-the-art baselines.

1 Introduction

Offline RL learns policies from fixed datasets, avoiding costly or dangerous exploration, but must address distributional shift and out-of-distribution actions. RAMBO addresses this challenge by adversarially modifying a learned model and reports stronger benchmark performance than existing state-of-the-art methods.

  • Offline RL learns policies from static datasets, avoiding the extensive exploration required by online RL.
  • Distributional shift makes conservatism essential because learned policies may execute state-action pairs outside the dataset distribution.
  • Model-based offline RL learns an environment model and trains policies with synthetic data, while prior methods use uncertainty-based reward penalties for conservatism.
  • RAMBO formulates offline RL as a zero-sum game and enforces conservatism by adversarially modifying learned transition dynamics.
  • RAMBO alternates policy optimisation with adversarial model optimisation without training a separate adversary policy.
  • RAMBO outperforms current state-of-the-art algorithms on D4RL benchmarks, while ablations identify adversarial model training as crucial to performance.

2 Related Work

Related work develops conservatism and robustness for offline or online RL, while RAMBO targets a practical model-based solution to a theoretically motivated maximin formulation.

  • Model-free offline RL uses policy constraints, value-function conservatism, importance sampling, or uncertainty quantification to address fixed-data learning.
  • Many model-based offline RL methods use maximum-likelihood models, whereas RAMBO produces pessimistic synthetic transitions for out-of-distribution behavior.
  • Prior work introduced the maximin formulation and PAC motivation for offline RL but did not propose a practical algorithm.
  • RAMBO proposes a practical model-based RL algorithm for solving that maximin formulation.
  • Robust MDP methods optimize worst-case performance over possible MDPs, typically assuming an uncertainty set specified in advance.
  • RAMBO adapts the two-player zero-sum RARL paradigm by adversarially modifying a learned model rather than training an adversary policy.

3 Preliminaries

These preliminaries define MDPs, offline data, learned model-based policy optimisation, and robust adversarial RL. They motivate RAMBO as a response to model exploitation caused by distributional shift.

  • MDPs and Offline RL: An MDP consists of state and action spaces, rewards, transition dynamics, an initial-state distribution, and a discount factor.
  • MDPs and Offline RL: A policy induces discounted state and state-action visitation distributions and an expected discounted return in an MDP.
  • MDPs and Offline RL: Offline RL seeks the best policy using only a fixed dataset of transitions from the MDP.
  • Model-Based Offline RL Algorithms: Model-based offline RL learns transition dynamics, forms an estimated MDP, and applies planning or RL to recover a policy.
  • Model-Based Offline RL Algorithms: Naive policy optimisation in a learned model can exploit inaccurate dynamics in state-action regions absent from the dataset.
  • Model-Based Offline RL Algorithms: MBPO trains an actor-critic value function on real and synthetic data generated by k-step rollouts from dataset states.
  • Robust Adversarial Reinforcement Learning: RARL formulates robust policy learning as a two-player zero-sum game and approximates it by alternating agent and adversary updates.

4 Problem Formulation

RAMBO addresses a maximin offline-RL formulation that constrains learned models near the dataset while leaving uncovered regions flexible. The formulation has PAC and pessimistic-value guarantees, and the paper develops a practical solver for it.

  • 4 Problem Formulation: The formulation assumes unknown transition and reward functions and represents them jointly through the learned model.
  • 4 Problem Formulation: The maximin problem seeks a policy that maximizes performance against an adversarial model constrained by dataset agreement.
  • 4 Problem Formulation: The constraint keeps candidate models close to the maximum-likelihood estimate on dataset state-action pairs.
  • 4 Problem Formulation: Because the constraint is averaged over the dataset distribution, it imposes no restriction on state-action regions not covered by the data.
  • 4.1 Theoretical Motivation: The PAC analysis guarantees that the resulting policy is approximately as good as policies whose state-action distributions are covered by the dataset.
  • 4.1 Theoretical Motivation: The performance gap is controlled by a discrepancy between the comparator policy's visitation distribution and the offline data distribution.
  • 4.1 Theoretical Motivation: The worst-case model's value function lower-bounds the true-environment value with high probability.
  • 4.1 Theoretical Motivation: The prior maximin work supplied theoretical motivation but no practical algorithm, motivating RAMBO's algorithmic development.

5 RAMBO-RL

RAMBO alternates policy optimisation with adversarial model updates to make the learned environment pessimistic while preserving dataset fidelity. Its model gradient modifies successor-state and reward likelihoods, and the final loss balances maximum-likelihood fitting against value reduction.

  • RAMBO alternates agent updates that increase expected value with adversarial model updates that decrease it.
  • The model-gradient method updates successor-state and reward likelihoods rather than policy action probabilities.
  • The adversarial model is constrained to remain accurate on dataset transitions while reducing the policy’s value in the learned model.
  • The final model loss adds an adversarial value term to the standard maximum-likelihood loss, with λ controlling their tradeoff.
  • RAMBO pretrains an MLE dynamics model, generates synthetic k-step rollouts, updates the policy with real and synthetic data, and then updates the model adversarially.

6 Experiments

RAMBO is evaluated across D4RL locomotion and maze benchmarks, with comparisons to model-based and model-free baselines, offline tuning, ablations, and a comparison to COMBO. It performs strongest on MuJoCo locomotion, while results are weaker for high-quality Medium-Expert data and AntMaze.

  • D4RL Performance: RAMBO achieves the best total score for the D4RL MuJoCo locomotion domains and performs especially well on Medium and Medium-Replay datasets.
  • D4RL Performance: RAMBO is outperformed by most baselines on Medium-Expert datasets, indicating lower suitability for high-quality datasets.
  • D4RL Performance: Model-based algorithms perform considerably worse than model-free approaches on AntMaze, although RAMBO scores above zero on most datasets.
  • Offline Tuning: RAMBOOFF shows slight performance degradation relative to online-tuned RAMBO but remains comparable to the best existing approaches on MuJoCo datasets.
  • Ablation of Adversarial Training: Removing adversarial updates degrades overall performance, although some datasets perform best without adversarial training.
  • Comparison to COMBO: RAMBO consistently outperforms COMBO on the Single Transition Example, where gradual pessimism makes policy optimisation less likely to get stuck in poor local maxima.

7 Conclusion and Future Directions

RAMBO is presented as a theoretically justified model-based offline RL approach that imposes conservatism by adversarially modifying learned transition dynamics. The paper reports state-of-the-art performance on standard benchmarks and identifies several future application directions.

  • RAMBO imposes conservatism by adversarially modifying the transition dynamics of a learned model.
  • The approach is theoretically justified and achieves state-of-the-art performance on standard benchmarks.
  • Future Directions: Future work includes applying RAMBO to image-space domains through latent-space dynamics perturbations.
  • Future Directions: The authors also propose using adversarially trained models for interpretability and extending the approach to online RL for robustness.

A Proof of Proposition 2

The appendix derives the model gradient by expanding the value function with Bellman’s equation, applying the product rule, and using visitation distributions and the log-derivative trick. It also explains that a state-action baseline can be subtracted without biasing the gradient estimate.

  • The proof begins by expressing the value function with Bellman’s equation and then applies the product rule.
  • The derivation defines φ(s →x, n) as the probability of reaching x from s after n policy steps in the learned MDP model.
  • φ(s) represents the improper discounted state visitation distribution under policy π and initial-state distribution µ0.
  • Applying the log-derivative trick yields the final model-gradient expression.
  • A quantity independent of the next state or reward can be subtracted, allowing Qπφ(s, a) to serve as a baseline without biasing the gradient estimate.

B.1 Model Training

RAMBO trains an ensemble of dynamics models with maximum likelihood and then updates them adversarially while jointly training the policy and value function on real and synthetic data. Performance depends strongly on rollout length and adversarial weighting, with RAMBO generally outperforming offline hyperparameter selection on MuJoCo domains.

  • Model Training: The model is an ensemble of neural networks that predicts next states and rewards with a Gaussian distribution conditioned on the current state and action.
  • Model Training: Seven dynamics models are pretrained by maximum likelihood, and the five with the lowest held-out validation error are retained for adversarial training.The held-out set contains 1000 dataset transitions.
  • Model Training: RAMBO updates all five retained models adversarially and randomly selects one model for each rollout.
  • Agent Training: The agent is trained with SAC using batches containing equal proportions of real and synthetic transitions.The real-data ratio is f = 0.5.
  • Hyperparameter Tuning: The main performance-sensitive hyperparameters are synthetic rollout length k and adversarial weighting λ; rollout lengths 2 or 5 worked well across most datasets.
  • Hyperparameter Tuning: λ must remain much smaller than 1 so maximum likelihood dominates on in-distribution samples, while excessive λ destabilizes training and overly small λ removes the adversarial effect.A value of 3e-4 generally obtained good performance.
  • Hyperparameter Tuning: RAMBOOFF selects hyperparameters offline using low and stable Q-values, whereas RAMBO selects the best-performing configuration across tested choices and seeds.
  • Results: RAMBOOFF achieves strong performance on MuJoCo domains, but RAMBO performs slightly better when online hyperparameter tuning is used.

C.1 Single Transition Example

The single-transition example compares RAMBO and COMBO under offline data coverage gaps. RAMBO gradually introduces pessimism and ultimately selects the optimal action covered by the dataset, whereas COMBO can become overly pessimistic and select a suboptimal covered action.

  • The one-dimensional agent executes one action from s0, receives reward equal to the successor state, and then terminates.
  • RAMBO initially overestimates out-of-distribution actions but regularises their Q-values after 50 iterations and selects the optimal dataset action, a ∈[0.7, 0.75].
  • COMBO with β = 0.2 is highly pessimistic about out-of-distribution actions at both 5 and 50 iterations, selecting a ∈[0.1, 0.15] instead of the optimal dataset action.
  • RAMBO introduces pessimism gradually, whereas COMBO applies it from the outset through its value-function update.
  • Across 20 seeds, RAMBO consistently outperforms COMBO, though its gradual pessimism may require more iterations to find a performant policy.

C.2 Results for each Hyperparameter Configuration

The experiments examine RAMBO across hyperparameter configurations and illustrate how adversarial transition training changes offline policy optimisation. The illustrative examples show accurate in-distribution modeling alongside pessimistic treatment of out-of-distribution actions, while behavior-cloning initialization provides a small performance improvement.

  • The reported configurations use three tested rollout-length and adversarial-weighting combinations, with highlighted values marking the best performance per dataset.
  • Table 6 compares three tested configurations of rollout length k and adversarial weighting λ using average normalized return over five seeds.
  • Naïve MBPO fits the dataset but overestimates values for actions outside the dataset, favoring an approximately unit action.
  • RAMBO preserves accurate transitions within the dataset while adversarial training generates pessimistic transitions for out-of-distribution actions.
  • Pessimistic synthetic transitions lead the illustrated SAC agent to select a ≈0.3, the best action within the dataset distribution.
  • Removing behavior-cloning initialization produces a small performance decrease for RAMBO on the MuJoCo locomotion experiments.
Loading 2204.12581v3…