Pessimistic Latent Task-aware Optimization for Robust Offline Meta-Reinforcement Learning

Chang-Hoon Jeong; Kisung Shin; Hanwool Sul; Byoung-Tak Zhang

cs.LGOct 11, 2026

Pessimistic Latent Task-aware Optimization for Robust Offline Meta-Reinforcement Learning

Authors: Chang-Hoon Jeong, Kisung Shin, Hanwool Sul, Byoung-Tak Zhang

Organizations: Dept. of Computer Science and Engineering, Seoul National University · Interdisciplinary Program in Artificial Intelligence, Seoul National University

Abstract

Context-based offline meta-reinforcement learning (COMRL) has advanced almost entirely through better task representations, which separate task identity from artifacts of the behavior policy so that the latent reflects the task alone. The implicit expectation has been that a clean representation will transfer to out-of-support (parameter-extrapolation) tasks. On such a task, however, the encoder produces a latent outside the training support, where the policy has never been trained. We address this gap with Pessimistic Latent Task-aware Optimization (PLATO), which exposes the policy to off-support latents during training. PLATO assumes a mild manifold hypothesis: task latents lie on a low-dimensional manifold, and out-of-support tasks lie farther from the training centroid. It exploits this geometry by perturbing inferred latents outward from the centroid to synthesize counterfactual tasks. A learned decoder ensemble then rolls out short trajectories at each perturbed latent, and the policy is updated on the imagined transitions under a pessimism weight that down-weights steps on which the ensemble disagrees. We prove that the perturbation leaves the training support by a controllable margin and that the out-of-support generalization gap is bounded by the same perturbation reach, and we give explicit conditions, verified empirically, under which the epistemic disagreement at the perturbed latent stays bounded away from zero. On eight MuJoCo continuous-control benchmarks, PLATO consistently improves out-of-support returns over strong COMRL baselines while remaining competitive on in-distribution tasks. Our code is available at https://github.com/deephoony/plato.

Figures & tables

Figure 1: Latent geometry on Ant-Goal. Background: decoder-ensemble disagreement

u(\mathbf{z})

over the latent space projected by principal component analysis (PCA). In-distribution (ID) tasks (blue) populate the low-

u

region around the centroid

c

(

+

). Out-of-support tasks (orange, labeled OOD) exhibit substantially higher disagreement on average, and our perturbations

\mathbf{z}\to\mathbf{z}^{\text{aug}}

(arrows,

\times

) reach the same high-

u

outward region, consistent with Assumption 1 , Proposition 1 , and Theorem 1 .

Figure 2: Conceptual overview of PLATO. From left to right: a latent $\mathbf{z}_{i}$ is inferred from the context of a training task; it is pushed outward from the training centroid $c$ to $\mathbf{z}_{i}^{\text{aug}}$ , off the training support $\mathcal{Z}_{\text{train}}$ ( Assumption 1 , Proposition 1 ); and the $M$ -member decoder ensemble, conditioned on $\mathbf{z}_{i}^{\text{aug}}$ , rolls out a short imagined trajectory from a real state $s_{0}$ , recording the per-step disagreement $u_{h}$ as its members drift apart. The policy $\pi_{\theta}$ is updated on real transitions and on the imagined ones, each imagined step weighted by $w_{h}=1/(1+\lambda_{\text{pess}}u_{h})$ , while the critic $Q_{\omega}$ is trained on real transitions only ( Theorem 2 ).

Environment	Offline PEARL	FOCAL	CSRO	UNICORN	ER-TRL	PLATO
Ant-Dir	-31 $\pm$ 33	68 $\pm$ 58	-17 $\pm$ 125	-49 $\pm$ 67	67 $\pm$ 48	139 $\pm$ 25
Ant-Goal	-491 $\pm$ 58	-480 $\pm$ 30	-527 $\pm$ 48	-517 $\pm$ 63	-489 $\pm$ 43	-472 $\pm$ 20
Cheetah-Vel	-230 $\pm$ 12	-199 $\pm$ 30	-188 $\pm$ 14	-203 $\pm$ 11	-201 $\pm$ 11	-179 $\pm$ 11
Humanoid-Dir	687 $\pm$ 58	499 $\pm$ 75	624 $\pm$ 83	580 $\pm$ 79	457 $\pm$ 45	806 $\pm$ 97
Hopper-Mass	402 $\pm$ 76	467 $\pm$ 43	494 $\pm$ 59	397 $\pm$ 55	444 $\pm$ 62	508 $\pm$ 26
Hopper-Friction	366 $\pm$ 31	334 $\pm$ 55	358 $\pm$ 48	338 $\pm$ 64	357 $\pm$ 21	369 $\pm$ 29

Table 1: Out-of-support returns on the held-out task splits. Mean

\pm

standard deviation over 5 seeds at the last epoch. Boldface marks the highest average return per row.

Environment	Offline PEARL	FOCAL	CSRO	UNICORN	ER-TRL	PLATO
Ant-Dir	183 $\pm$ 39	589 $\pm$ 20	547 $\pm$ 49	321 $\pm$ 45	605 $\pm$ 34	617 $\pm$ 29
Ant-Goal	-242 $\pm$ 15	-218 $\pm$ 10	-214 $\pm$ 13	-249 $\pm$ 12	-207 $\pm$ 11	-215 $\pm$ 10
Cheetah-Vel	-73 $\pm$ 16	-69 $\pm$ 16	-68 $\pm$ 7	-64 $\pm$ 4	-64 $\pm$ 6	-61 $\pm$ 11
Humanoid-Dir	688 $\pm$ 74	550 $\pm$ 44	632 $\pm$ 47	616 $\pm$ 46	529 $\pm$ 61	858 $\pm$ 43
Hopper-Mass	524 $\pm$ 61	503 $\pm$ 59	556 $\pm$ 31	517 $\pm$ 53	502 $\pm$ 53	584 $\pm$ 12
Hopper-Friction	517 $\pm$ 46	526 $\pm$ 68	584 $\pm$ 15	474 $\pm$ 101	521 $\pm$ 44	526 $\pm$ 42

Table 2: In-distribution performance on MuJoCo benchmarks. Same format as Table 1 . Per-environment task ranges and ID/out-of-support split definitions are in Appendix E .

Figure 3: Off-support disagreement stays large. Per-task ensemble disagreement at the trained PLATO checkpoint. ID values (blue) are near zero, whereas out-of-support values (red, labeled OOD) are several-fold to an order of magnitude larger, consistent with Theorem 1 .

Figure 4: Component ablations of PLATO on out-of-support returns. Learning curves on six MuJoCo environments. default is the full algorithm; no-direction replaces the centroid-outward direction

\hat{n}_{i}

with a random unit vector; no-weight sets

\lambda_{\text{pess}}=0

; Q-synth also trains the critic on imagined Bellman targets (the variant analyzed by Theorem 2 ). Curves show mean over 5 seeds; shaded regions indicate one standard deviation.

Appendix figures & tables9 assets

Supplementary material from the paper’s appendix.

Appendix

Environment	Type	$\|S\|$	$\|A\|$	$\|\Theta\|$	ID	Out-of-support
Ant-Dir	Reward shift	27	8	30	20	10
Ant-Goal	Reward shift	27	8	30	20	10
Cheetah-Vel	Reward shift	17	6	30	20	10
Humanoid-Dir	Reward shift	376	17	30	20	10
Hopper-Mass	Dynamics shift	11	3	30	20	10
Hopper-Friction	Dynamics shift	11	3	30	20	10

Appendix

Table 3: Environment summary.

|S|

and

|A|

are the state and action dimensions, and

|\Theta|

is the number of tasks used in this paper. The ID and out-of-support columns give the numbers of in-distribution training tasks and out-of-support evaluation tasks; the out-of-support split is held out far from the training support (see the per-environment descriptions below). For MetaWorld Reach, the remaining

10

tasks (the corners of the goal box) are used only for the extrapolation-limit analysis in Appendix K .

Figure 5: Ensemble disagreement tracks task difficulty; latent distance does not. Each point is one Ant-Dir task (

20

training and

10

out-of-support tasks). Vertical axis: normalized regret of a naively trained agent (CSRO),

(R_{\text{expert}}-R_{\text{CSRO}})/|R_{\text{expert}}|

. Left: distance of the task’s mean latent from the training centroid. Right: decoder-ensemble disagreement

u

at the task’s latent (log scale).

\rho

is the Spearman rank correlation over all

30

tasks.

Table 9

Figure 6: Out-of-support sensitivity analysis. Each row sweeps one hyperparameter, with the other two held at their middle swept value. Top: perturbation magnitude

\alpha\in\{0.1,0.7,1.5\}

. Middle: pessimism strength

\lambda_{\text{pess}}\in\{0,5,10\}

. Bottom: warm-up

T_{\text{warm}}\in\{5\text{K},10\text{K},30\text{K}\}

. Columns: Ant-Goal and Cheetah-Vel (reward-shift), Hopper-Mass and Walker-Friction (dynamics-shift). Curves show mean out-of-support return over 5 seeds.

$M$	Cheetah-Vel	Hopper-Mass	Walker-Friction
$3$	$-186\pm 12$	$479\pm 36$	$564\pm 145$
$5$	$-187\pm 2$	$478\pm 8$	$655\pm 53$
$10$	$-179\pm 11$	$508\pm 26$	$649\pm 27$

Appendix

Table 6: Out-of-support return for different ensemble sizes

M

M=10

is the main configuration ( Table 1 ,

5

seeds);

M\in\{3,5\}

are averaged over

3

seeds.

Figure 7: Learning curves on the eight MuJoCo OMRL benchmarks, in-distribution. Each panel shows normalized return over training, averaged across 5 seeds; shaded regions indicate one standard deviation.

Figure 8: Learning curves on the eight MuJoCo OMRL benchmarks, out-of-support. Same format as Figure 7 .

Figure 9: MetaWorld Reach, out-of-support tasks. Success rate (left) and average return (right) over training, averaged over

3

seeds; shaded regions indicate one standard deviation.

	In-distribution		Out-of-support
Environment	MerPO	PLATO	MerPO	PLATO
Ant-Dir	678 $\pm$ 34	617 $\pm$ 29	70 $\pm$ 18	139 $\pm$ 25
Ant-Goal	-225 $\pm$ 12	-215 $\pm$ 10	-481 $\pm$ 20	-472 $\pm$ 20
Humanoid-Dir	652 $\pm$ 51	858 $\pm$ 43	577 $\pm$ 68	806 $\pm$ 97
Hopper-Mass	602 $\pm$ 19	584 $\pm$ 12	516 $\pm$ 24	508 $\pm$ 26
Hopper-Friction	576 $\pm$ 33	526 $\pm$ 42	329 $\pm$ 27	369 $\pm$ 29

Appendix

Table 7: Comparison with MerPO. Returns on the in-distribution and out-of-support splits, mean

\pm

standard deviation over 6 seeds at the last epoch for MerPO; the PLATO values are those of Tables 1 and 2 . Boldface marks the higher average return in each pair.

Meta-Reinforcement Learning

Meta-reinforcement learning trains agents across tasks so they can adapt quickly to new environments or objectives.

See all 40 papers

Mar 3, 2026

Temporal Consistency Improves Generalization in Contextual Offline Meta Reinforcement Learning

Offline meta-reinforcement learning seeks to learn a policy that generalizes to new related tasks online. Context-based methods infer a task representation from transition histories, yet learning an effective task representation without supervision remains challenging. Existing methods relying on contrastive learning learn discriminative task representations, but fail to identify task-specific dynamics, while relying on reconstruction can be insufficient to model long-horizon dependencies, limiting generalization to new tasks. We investigate the impact of temporal consistency in latent space on task representation learning, showing that enforcing multi-step predictions in latent space encourages task representations that are able to capture task-dependent dynamics while preventing representation collapse. We provide theoretical analysis characterizing sources of error in value estimation and show through extensive experiments on MuJoCo, Contextual DeepMind Control, and MetaWorld benchmarks that temporal consistency significantly improves both zero-shot and few-shot generalization.

Mohammadreza Nakheai, Aidan Scannell, Kevin Luck +1

Aalto University · Toyota Research Institute · Vrije Universiteit Amsterdam

Representation Learning Offline RL Representation Learning for RL Meta-Reinforcement Learning Generalization in RL Temporal Consistency+2

Ask PDF

May 30, 2026

Behavior-Invariant Task Representation Learning with Transformer-based World Models for Offline Meta-Reinforcement Learning

Offline meta-reinforcement learning leverages static datasets to enable agents to generalize to unseen environments by combining offline efficiency with meta-learning adaptability, yet it faces key challenges from context and policy distribution shifts. These issues hinder agents from adapting to online environments, and are further exacerbated under sparse-reward settings. As a result, agents often become trapped in an inherent pattern dilemma, failing to achieve robust generalization. In this work, we propose a novel framework that integrates information-theoretic task representation learning with a Transformer-based stochastic world model. Our approach extracts task-defining latent variables that are invariant to behavior policy, thereby effectively mitigating the context distribution shift. To further handle policy shift and model exploitation, we apply a conservative value penalty to imagination-based rollouts, preventing the policy from exploiting model inaccuracies while maintaining robust adaptation. Extensive evaluations demonstrate that our method outperforms state-of-the-art approaches, with superior stability and generalization under out-of-distribution and sparse-reward settings.

Fuyuan Qian, Menglong Zhang, Song Wang +1

Department of Biomedical Engineering, Southern University of Science and Technology, Shenzhen, Guangdong, China · AI Thrust, The Hong Kong University of Science and Technology (Guangzhou), Guangzhou, Guangdong, China · Omni-Intelligence, Shenzhen, Guangdong, China +1

Representation Learning Offline RL Meta-Reinforcement Learning Model-Based RL

Ask PDF

Dec 24, 2025

Meta-RL with Bayesian Linear Task Models

Deep Bayesian reinforcement learning adapts to unseen tasks by inferring latent transition and reward models, but existing methods typically rely on variational posteriors and evidence lower bounds, introducing approximation error and unstable task representations. We introduce GLiBRL, a deep Bayesian RL framework that combines generalised linear task models with learnable non-linear basis functions. GLiBRL features conjugate Bayesian inference, yielding exact, sequential posterior updates over task parameters and model noise, together with a closed-form marginal likelihood that eliminates variational inference. The update is naturally permutation-invariant, allowing GLiBRL to integrate with both off- and on-policy algorithms. GLiBRL also learns task representation admitting an exact kernel identity, relating distances between task representations to kernel discrepancies over the task contexts. Compared against eight representative or recent meta reinforcement learning methods, GLiBRL achieves the highest aggregate zero-shot test performance on both the MuJoCo locomotion and MetaWorld manipulation benchmarks.

Jingyang You, Hanna Kurniawati

School of Computing Australian National University

Representation Learning Reinforcement Learning Bayesian RL Meta-Reinforcement Learning Generalization in RL Model-Based RL

Ask PDF

Abstract

Meta-Reinforcement Learning

Explore similar work

Temporal Consistency Improves Generalization in Contextual Offline Meta Reinforcement Learning

Behavior-Invariant Task Representation Learning with Transformer-based World Models for Offline Meta-Reinforcement Learning

Meta-RL with Bayesian Linear Task Models