Post-Hoc Robustness for Model-Based Reinforcement Learning

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.

• The paper proposes a post-hoc robust MPC method that enhances RL inference by adversarially perturbing the learned dynamics within a bounded uncertainty set.
• The paper employs adversarial rollouts and ensemble-based OOD truncation to optimize worst-case performance while mitigating error accumulation.
• The paper demonstrates improved robustness on MuJoCo benchmarks, effectively balancing nominal performance with safety under environmental perturbations.

Post-Hoc Robustness for Model-Based Reinforcement Learning

Problem Statement and Motivation

Robustness in RL, particularly in sim2real and sim2sim transfer scenarios, remains an enduring challenge due to environmental discrepancies between training and deployment. Classic solutions such as domain randomization optimize for expectation, not worst-case, often exposing agents to brittle deployments under rare system perturbations. Robust MDP (RMDP) formulations instead target adversarial, worst-case transitions, typically operationalized via adversarial training with sets of perturbed transition dynamics and regularized policy optimization.

This work departs from the standard paradigm by addressing the problem of post-hoc robustness—improving the robustness of a previously trained model-based RL (MBRL) agent at inference time without any additional data collection or re-training. The ability to decouple robustification from training is essential for practical deployments, where re-training is computationally costly or infeasible, and provides flexibility in defining robustness tradeoffs after initial model deployment.

Figure 1: The scope is post-hoc robustification, operating solely at inference without changing or fine-tuning the original training.

Methodology

The proposed approach instantiates post-hoc robust policy improvement as a model-predictive control (MPC) procedure applied to the learned transition model $p_\theta$ and nominal value function $V^\pi$. At each inference step, rather than outputting the nominal action, the method performs adversarial rollouts on the learned model to locally optimize for policy robustness within a bounded uncertainty set.

Key elements include:

• Adversarial Rollouts: At every timestep of the MPC horizon, adversarial PGD is used to perturb the next-state prediction to the worst-case within an $L_2$ ball, as formalized by an Optimal Transport Cost (OTC) uncertainty set. The adversary seeks to minimize the value over the nominal value function, yielding pessimistic rollouts.
• MPC over Robustified Rollouts: Actions are obtained by running a sampling-based MPPI algorithm initialized with the learned policy, but optimizing over rollouts where transitions have been adversarially perturbed.
• Out-of-Distribution Safeguards: Due to recursive model usage and adversarial perturbations, OOD errors quickly accumulate. Inspired by offline RL (MOReL), the method employs an ensemble discrepancy-based rollout truncation. If epistemic uncertainty, measured as the discrepancy between model ensemble predictions, exceeds a threshold, rollouts are aborted, preventing exploitation of states outside the reliable training distribution.

Empirical Evaluation

Evaluation Protocol

Robustness is assessed under both single and simultaneous parameter perturbations in Gymnasium MuJoCo Reacher-v4 and Hopper-v4 benchmarks. The baseline is unmodified MBPO, with non-robust MPC and the proposed robust MPC methods compared through a suite of controlled ablation studies.

Robustness Under Perturbation

Results clearly demonstrate that robust MPC consistently decreases performance degradation under increasingly severe environmental parameter disturbances.

Figure 2: Robust MPC yields higher returns compared to MBPO and non-robust MPC under parameter perturbations on Reacher-v4.

Figure 3: Visualization of the Reacher-v4 environment.

Simultaneous Perturbation Robustness

When evaluated under the joint effect of pairs of perturbations (e.g., damping and gear in Reacher, mass and friction in Hopper), robust MPC provides measurable improvement in maintaining task performance over the baseline and non-robust variants.

Influence of Rollout Depth and Uncertainty Set Radius

Ablation results show that increasing the MPC rollout horizon further boosts robustness up to an optimal point; excessive depth degrades performance due to compounding model error.

Figure 4: Performance as a function of increasing MPC rollout depth reveals an optimal trade-off between robustness and model error on Hopper-v4.

Larger uncertainty set radii initially improve worst-case performance but eventually cause collapse due to excessive OOD queries as the perturbation exceeds trustworthy model support.

Figure 5: Effect of uncertainty set radius on robustness, with an identified stable regime before performance deteriorates.

Out-of-Distribution Mitigation

Ablations indicate that truncating rollouts based on ensemble disagreement is crucial; omitting this leads to inferior performance even relative to the vanilla baseline, confirming the practical necessity of OOD mitigation.

Computational Costs

Across modern GPUs, the robust MPC inference (~0.09s/step on RTX 4090 for Hopper-v4, $N$=20) is computationally feasible, confirming the method’s applicability for inference-time deployment.

Theoretical and Practical Implications

By decoupling robustification from training, this methodology allows specifying robustness requirements after deployment and adapting the effective trade-off between nominal performance and robustness. The use of adversarial rollouts on a learned world model enables practical post-hoc robustification across arbitrary environments without necessitating further environment interactions.

There are, however, limitations. The method is pessimistic with respect to the nominal value function rather than a true robust value function, particularly for $N>1$ rollouts. While 1-step rollouts coincide with robust Bellman updates, multi-step adversarial rollouts do not come with robust optimality guarantees. Nonetheless, empirical outcomes on standard control benchmarks validate their practical utility.

The addition of strong OOD truncation addresses the primary pitfall of model exploitation but does not guarantee safety in broader OOD scenarios. Extending this method to high-dimensional, latent-space models (e.g., vision-based control), continuous adaptation, or stochastic disturbances could further enhance real-world applicability.

Relation to Prior Work

The method builds upon literature on adversarial robust RL, RMDPs, and robust MPC, but critically operates at post-hoc inference with no retraining [see (2606.03521)]. Unlike context-adaptive RL methods that require explicit context features during both training and deployment, this approach modifies agent action selection mechanically using the predicted dynamics model, making it suitable where context is unobserved or intractable to enumerate.

Conclusion

This work presents a methodologically rigorous and computationally practical approach for realizing post-hoc robustness in MBRL agents via adversarial model-predictive rollouts and OOD truncation. Empirical studies in MuJoCo control environments confirm substantial improvements in robustness under parameter perturbations, with ablations affirming the contribution of each design choice. Future research should explore scaling to more complex domains and bridging the theoretical gap for multi-step pessimistic planning.