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

Temporal Consistency Improves Generalization in Contextual Offline Meta Reinforcement Learning
Mar 3, 2026Mohammadreza Nakheai, Aidan Scannell, Kevin Luck +1Representation Learning Offline RL
Aalto University · Toyota Research Institute · Vrije Universiteit Amsterdam
Behavior-Invariant Task Representation Learning with Transformer-based World Models for Offline Meta-Reinforcement Learning
May 30, 2026Fuyuan Qian, Menglong Zhang, Song Wang +1Representation Learning Offline RL
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
Meta-RL with Bayesian Linear Task Models
Dec 24, 2025Jingyang You, Hanna KurniawatiRepresentation Learning Reinforcement Learning
School of Computing Australian National University

Mar 3, 2026J/K move · Enter open · S save

Temporal Consistency Improves Generalization in Contextual Offline Meta Reinforcement Learning

Mohammadreza Nakheai, Aidan Scannell, Kevin Luck, Joni Pajarinen

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

Open paper AskPDF arXiv

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.

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