Post-Hoc Robustness for Model-Based Reinforcement Learning

2026-06-02

AI Summary

This work addresses the insufficient robustness of deep reinforcement learning agents under adversarial environmental perturbations by proposing a post-hoc robustification method that requires no retraining. During inference, the approach integrates a learned environment model with a nominal policy and refines the policy through adversarial receding-horizon optimization within a model predictive control framework. Perturbations are characterized via a bounded uncertainty set, and the resulting optimization problem is approximately solved using projected gradient descent. As the first method to achieve post-hoc robustification for deep reinforcement learning agents, it shifts the adversarial objective from the true environment to the learned model, thereby mitigating out-of-distribution generalization issues. Evaluated on perturbed Gymnasium MuJoCo tasks, the method yields substantial performance improvements while adhering to real-time inference computational constraints.

Abstract

To improve the real-world applicability of reinforcement learning (RL), the field of adversarially robust RL studies how to train agents under adversarial environment perturbations. In this setting, a protagonist agent optimizes a policy under environmental perturbations from an adversary, resulting in a zero-sum Markov game. When adversarially robust RL is combined with model-based RL, the adversary can target a learned transition model instead of the training environment. Extending this idea, this work introduces post-hoc robustification of deep RL agents at inference time. By using the learned model in combination with a trained nominal policy, our approach performs a robust policy improvement step. The goal is to improve robustness without any additional training of neural networks. Specifically, we utilize model-predictive control under adversarial rollouts, which are approximated via projected gradient descent within a bounded uncertainty set. Furthermore, these offline rollouts are performed while considering and mitigating out-of-distribution issues. The proposed methodology is validated by demonstrating significant improvements in robustness when the algorithm is evaluated in perturbed Gymnasium MuJoCo environments, while considering the computational limitations of the post-hoc inference setting.

Problem

• post-hoc robustness
• model-based reinforcement learning
• adversarial perturbations
• robust policy improvement
• inference-time adaptation

Innovation

• post-hoc robustification
• model-based reinforcement learning
• adversarial rollouts
• model-predictive control
• out-of-distribution mitigation