Reinforcement Learning (RL) agents often struggle to generalize knowledge to new tasks, even those structurally similar to ones they have mastered. Although recent approaches have attempted to mitigate this issue via zero-shot transfer, they are often constrained by predefined, discrete class systems, limiting their adaptability to novel or compositional task variations. We propose a significantly more generalized approach, replacing discrete latent variables with natural language conditioning via a text-conditioned Variational Autoencoder (VAE). Our core innovation utilizes a Large Language Model (LLM) as a dynamic semantic operator at test time. Rather than relying on rigid rules, our agent queries the LLM to semantically remap the description of the current observation to align with the source task. This source-aligned caption conditions the VAE to generate an imagined state compatible with the agent's original training, enabling direct policy reuse. By harnessing the flexible reasoning capabilities of LLMs, our approach achieves zero-shot transfer across a broad spectrum of complex and truly novel analogous tasks, moving beyond the limitations of fixed category mappings. Code and videos are available here.
Figures & tables
Figure 1: The overall idea of ASPECT. The agent uses an LLM to semantically remap the target observation (e.g., target task: “pick yellow duckie from blue room”) to a source-aligned description (e.g., source task: “pick blue box from grey room”), enabling the direct application of the pre-trained source policy to the target task.
Figure 2: Visualizations of the three environments used in our experiments: (a) MiniGrid (2D navigation), (b) MiniWorld (3D egocentric), and (c) Manipulation (continuous control).
Case
Env.
Target Task Description
1. Unseen Objects (Different reward object)
MiniGrid
Pick purple box , avoid green ball .
MiniWorld
Pick yellow duckie , avoid green ball .
Custom
Pick heavy circle & light square (inverted weights).
2. Combined Shift (Visual + Semantic)
MiniGrid
Blue Walls. Pick purple box .
MiniWorld
Wood/Brick textures. Pick yellow duckie .
3. Reversed Task (Prior as Distractor)
MiniGrid
Pick purple box , avoid red ball (source target).
Table 1: Target Task Settings (Evaluation). The agent must generalize zero-shot to these novel scenarios.
MiniWorld
MiniGrid
Manip.
Method
Tgt
Dstr
Tgt
Dstr
Lift
Brk
Source
2.80 ± 0.80
0.20 ± 0.20
0.00 ± 0.00
0.00 ± 0.00
2.80 ± 1.06
4.60 ± 0.04
FT (Conv.)
10.0 ± 0.00
0.60 ± 0.40
9.60 ± 0.24
0.00 ± 0.00
9.80 ± 0.20
0.00 ± 0.00
FT (Lim.)
8.00 ± 1.04
0.40 ± 0.24
7.20 ± 0.86
0.40 ± 0.40
6.60 ± 0.87
2.40 ± 1.51
SF-Simp.
0.20 ± 0.20
0.00 ± 0.00
5.60 ± 1.28
0.00 ± 0.00
-
-
SF-Rec.
0.00 ± 0.00
0.00 ± 0.00
9.00 ± 0.20
0.00 ± 0.00
-
-
Table 2: Consolidated Experimental Results for Case 1 (Generalization to Unseen Objects). Metrics are reported as counts (out of 10) ± std dev. ”Tgt” = Target, ”Dstr” = Distractor, ”Lift” = Successful Lifts, ”Brk” = Broken Objects.
MiniWorld
MiniGrid
Method
Tgt
Dstr
Tgt
Dstr
Source
3.20 ± 0.66
0.20 ± 0.20
0.00 ± 0.00
0.00 ± 0.00
FT (Conv.)
9.20 ± 0.37
0.40 ± 0.24
9.80 ± 0.20
0.00 ± 0.00
FT (Lim.)
7.60 ± 0.50
0.40 ± 0.24
6.20 ± 0.73
0.60 ± 0.24
SF-Simp.
0.00 ± 0.00
0.00 ± 0.00
4.80 ± 1.62
0.00 ± 0.00
SF-Rec.
0.00 ± 0.00
0.00 ± 0.00
4.60 ± 2.09
0.00 ± 0.00
Table 3: Consolidated Experimental Results for Case 2 (Combined Generalization). Metrics are reported as counts (out of 10) ± std dev. ”Tgt” = Target, ”Dstr” = Distractor.
MiniWorld
MiniGrid
Method
Target
Old Target
Target
Old Target
Source
2.80 ± 0.66
8.20 ± 0.37
0.00 ± 0.00
7.20 ± 0.86
FT (Conv.)
9.80 ± 0.24
0.00 ± 0.00
9.80 ± 0.20
0.00 ± 0.00
FT (Lim.)
5.60 ± 1.36
0.40 ± 0.24
3.40 ± 0.87
0.80 ± 0.37
SF-Simple
0.00 ± 0.00
0.00 ± 0.00
3.80 ± 1.68
0.00 ± 0.00
SF-Reconstruction
0.00 ± 0.00
0.00 ± 0.00
1.20 ± 0.96
0.00 ± 0.00
Table 4: Consolidated Experimental Results for Case 3 (Reverse Task). Metrics are reported as counts (out of 10) ± std dev. ”Old Target” refers to the object rewarding in source but distracting in target.
Appendix figures & tables9 assets
Supplementary material from the paper’s appendix.
Appendix
Figure 3: PPO Fine-tuning in MiniWorld in each case.
Case
Rewarding Object Picked
Distractor Object Picked
Case 1
8.60 ± 0.50
0.00 ± 0.00
Case 2
8.40 ± 0.40
0.00 ± 0.00
Case 3
8.60 ± 0.40
0.00 ± 0.00
Appendix
Table 5: Experimental results of ASPECT for MiniWorld environment using unstructured, VLM-generated captions. The method demonstrates robustness to noise and variability in the conditioning text.
Figure 4: Fine-tuning curves for MiniGrid and Fragile Object Manipulation.
Figure 5: Learning curve for Successor Features (SF) in MiniGrid. The agent interacts with environments sequentially, with performance dips indicating task transfers.
Figure 6: Learning curve for Successor Features (SF) in MiniWorld. The flat reward curve indicates a failure to learn or transfer to the target tasks.
Figure 7: Visualization of spatial latent channels. The first two rows of the latent feature maps remain consistent between (b) and (d), capturing the preserved background structure, while other channels shift to reflect the semantic change from (a) to (c).
Figure 8: Text-guided image generation demonstrating latent disentanglement. Structural features captured in z remain fixed (preserving the background layout) while the text modifies the semantic appearance. Left: MiniWorld, Right: MiniGrid.
Figure 9: Comparison of real target observations and the corresponding source-aligned imagined states generated by ASPECT. We display 10 examples demonstrating the robust visual remapping capability.
Figure 10: Visualisation of failure cases. The model struggles with extreme close-ups, leading to generation artefacts, and occasionally fails to generate the source object even when the target is visible.
Humans excel at analogical reasoning - applying knowledge from one task to a related one with minimal relearning. In contrast, reinforcement learning (RL) agents typically require extensive retraining even when new tasks share structural similarities with previously learned ones. In this work, we propose MAGIK, a novel framework that enables RL agents to transfer knowledge to analogous tasks without interacting with the target environment. Our approach leverages an imagination mechanism to map entities in the target task to their analogues in the source domain, allowing the agent to reuse its original policy. Experiments on custom MiniGrid and MuJoCo tasks show that MAGIK achieves effective zero-shot transfer using only a small number of human-labelled examples. We compare our approach to related baselines and highlight how it offers a novel and effective mechanism for knowledge transfer via imagination-based analogy mapping.
Ajsal Shereef Palattuparambil, Thommen George Karimpanal, Santu Rana
Applied Artificial Intelligence Initiative, Deakin University · School of IT, Deakin University
Post-training Vision-Language-Action (VLA) models via reinforcement learning (RL) in learned world models has emerged as an effective strategy to adapt to new tasks without costly real-world interactions. However, while using imagined trajectories reduces the sample complexity of policy training, existing methods still heavily rely on task-specific data to fine-tune both the world and reward models, fundamentally limiting their scalability to unseen tasks. To overcome this, we argue that world and reward models should capture transferable physical priors that enable zero-shot inference. We propose RAW-Dream (Reinforcing VLAs in task-Agnostic World Dreams), a new paradigm that completely disentangles world model learning from downstream task dependencies. RAW-Dream utilizes a world model pre-trained on diverse task-free behaviors for predicting future rollouts, and an off-the-shelf Vision-Language Model (VLM) for reward generation. Because both components are task-agnostic, VLAs can be readily finetuned for any new task entirely within this zero-shot imagination. Furthermore, to mitigate world model hallucinations, we introduce a dual-noise verification mechanism to filter out unreliable rollouts. Extensive experiments across simulation and real-world settings demonstrate consistent performance gains, proving that generalized physical priors can effectively substitute for costly task-dependent data, offering a highly scalable roadmap for VLA adaptation.
Yucen Wang, Rui Yu, Fengming Zhang +5
Nanjing University · Interns at Microsoft · University of Illinois at Urbana-Champaign +4
Compositional generalization is essential for reaching unseen goals under novel contextual variations in offline goal-conditioned reinforcement learning (GCRL), where a generalist goal-reaching agent must be learned from limited data. Most prior approaches pursue this via trajectory stitching over temporally contiguous segments, which limits composing behaviors across varying contexts. To overcome this limitation, we formalize analogy transduction as synthesizing new plans by composing task-endogenous analogies with given contexts and propose a novel analogy representation tailored for it. Grounded in our theory, this analogy representation captures what changes under optimal task execution, remains invariant to contextual variations, and is sufficient for optimal goal reaching. We further contend that generalization to unseen analogy-context pairs is a practical obstacle in analogy transduction, and introduce a new approach for offline GCRL that enables analogy transduction beyond seen pairs to unseen combinations. We empirically demonstrate the effectiveness of our approach on OGBench manipulation environments, substantially outperforming prior methods that do not perform analogy transduction. Project page: https://rllab-snu.github.io/projects/CTA/
Junseok Kim, Dohyeong Kim, Mineui Hong +1
Department of Electrical and Computer Engineering and ASRI, Seoul National University · Independent researcher · Robotics Institute, Carnegie Mellon University