Improving Policy Learning via Language-Guided State Representation in World Models

Zhang, Lixuan; Kan, Meina; Shan, Shiguang; Chen, Xilin

Abstract

World models have emerged as powerful technology for facilitating policy learning in robotics by providing predictive representations of environmental dynamics. A crucial component of such models is the internal state representation, which serves as a bridge between observation, decision-making, and future state prediction. However, in multi-task settings, the presence of task-irrelevant information within the representation, especially under limited information capacity, can significantly impair the effectiveness of policy learning.

To address this challenge, we propose Language-Guided State Representation (LGSR), a novel method that leverages language-based task descriptions as priors to guide the encoding of state representations within the world model. By incorporating language instructions, LGSR promotes the extraction of task-relevant features while suppressing irrelevant information, thereby enhancing the alignment between representation learning and task objectives.

We evaluate LGSR on the CALVIN benchmark, a suite of language-conditioned robotic manipulation tasks. Experimental results demonstrate that incorporating language priors into the state representation enhances the extraction of task-relevant information, thereby significantly improving the efficiency and success rate of downstream policy learning.

Method

Let the task decide what the state should remember.

LGSR injects language at the encoder stage, compressing visual observations into a task-conditioned state before the world model predicts the future and the policy predicts actions.

LGSR architecture: CLIP language and MAE visual features enter a multimodal perceiver encoder, whose task-conditioned state feeds a shared world model and policy transformer. — The LGSR framework. Frozen language and vision encoders feed a Multimodal Perceiver Encoder (MPE); the resulting compact state supports future prediction and action learning.

Results

Strong long-horizon control with a compact trainable model.

4.22CALVIN Avg. Len.Best reported LGSR variant
69MTrainable parametersNo robot-data pretraining
40 HzInference frequencyMeasured on RTX 4090
70%Training dataMatches full-data LGSR performance

Long-horizon multi-task language control

Models train on environments A, B, and C, then execute five-instruction chains in unseen environment D. LGSR with the proxy target completes 4.22 subtasks on average, exceeding the scratch Seer baseline by 0.58 and OpenVLA by 0.95.

Selected CALVIN ABC-D results from the final manuscript.
Method	1	2	3	4	5	Avg. Len.	Trainable
OpenVLA	91.3%	77.8%	62.0%	52.1%	43.5%	3.27	~100M
Seer (scratch)	93.0%	82.4%	73.0%	62.6%	53.3%	3.64	65M
LGSR	95.1%	88.4%	81.1%	73.2%	66.0%	4.04	65M
LGSR + proxy	96.3%	90.8%	84.5%	78.1%	71.9%	4.22	69M

Average sequence length versus CALVIN training-data proportion for LGSR, its language-free variant, and Seer. — **Data efficiency.** At 70% of the language-annotated split, LGSR already exceeds Seer trained with the full annotated and unannotated sets.

Average sequence length versus state representation token count for LGSR and the variant without language guidance. — **Capacity robustness.** Language guidance helps from one to six state tokens; with one token, LGSR delivers a 17.7% relative improvement.

Effect of masking Goal, Constraint, Action, Spatial, Logic, and Other instruction components. — **What words matter?** Goal, constraint, and spatial words are most important for learning task-relevant states; action words contribute little, while removing distracting logic or filler can help.

Which proxy target works?

Proxy-target ablation on CALVIN ABC-D.
Target	1	2	3	4	5	Avg.
Binary successrₜ = 1[t = T]	92.4%	83.5%	74.6%	66.5%	58.6%	3.76
Elapsed progressrₜ = t/T	94.8%	87.0%	78.8%	70.2%	62.8%	3.93
Time-to-gorₜ = (T − t)/T	96.3%	90.8%	84.5%	78.1%	71.9%	4.22

Is task identity in the state?

Task-ID prediction accuracy from frozen representations; mean ± sample standard deviation over 20 train/test splits.
Representation	CALVIN	LIBERO-Long
Seer	48.54 ± 1.19%	99.62 ± 0.15%
LGSR	97.94 ± 0.32%	100.00 ± 0.00%
LGSR + proxy	98.38 ± 0.41%	100.00 ± 0.00%

Language fusion contributes the large separability gain. The much smaller proxy effect depends on the feature probe and benchmark.

Representation Analysis

What survives a one-token language-guided bottleneck? An inverse decoder lifts the frozen one-token LGSR back into MAE feature space before the same frozen MAE decoder reconstructs the scene.

Raw view for lifting the blue block from the drawer. — Lift the **blue block** from the drawer.

MAE reconstruction for lifting the blue block from the drawer. — Lift the **blue block** from the drawer.

Raw view for lifting the pink block from the drawer. — Lift the **pink block** from the drawer.

MAE reconstruction for lifting the pink block from the drawer. — Lift the **pink block** from the drawer.

Raw view for lifting the red block from the drawer. — Lift the **red block** from the drawer.

MAE reconstruction for lifting the red block from the drawer. — Lift the **red block** from the drawer.

Raw view for closing the drawer. — **Close the drawer.**

MAE reconstruction for closing the drawer. — **Close the drawer.**

Fine appearance details are discarded, but the language-specified block—or the drawer itself—remains. This directly supports the claim that MPE spends limited state capacity on task-relevant scene content.

Additional Material and Analyses

Additional implementation details, CALVIN analyses, cross-benchmark experiments, diagnostics, and limitations that complement the main paper.

Architecture

Transformer Inputs and Prediction Slots

At timestep i, the transformer receives a context block Aᵢ = [language, proprioception, visual state] followed by prediction slots Bᵢ = [mₒ, mₐ]. Ground-truth future observations and actions are decoding targets, never input tokens. World-model and policy features are produced by one shared forward pass.

K = 2 attention mask. ● is allowed; — is masked. Rows are queries and columns are keys.
Query ↓ / Key →	A₁	mₒ¹	mₐ¹	A₂	mₒ²	mₐ²
A₁	●	—	—	—	—	—
mₒ¹	●	—	—	—	—	—
mₐ¹	●	●	—	—	—	—
A₂	●	—	—	●	—	—
mₒ²	●	—	—	●	—	—
mₐ²	●	—	—	●	●	—

Causal over time. Step i cannot read any token from a later timestep.
Prediction slots are sinks. Context and future slots cannot read prediction tokens, preventing target leakage.
One deliberate exception. Same-step action token mₐ can read mₒ, so the policy reads the internally predicted future feature.

Inference: use the same mask tokens, decode the latest action chunk, execute its first action, advance history by one observation, and repeat. The future decoder is an auxiliary training signal; control is produced from the action feature.

Additional CALVIN Experiments

Prediction-Target Ablation

All rows use the same LGSR pipeline; only the future-vision target changes. Zero-mean patch prediction removes low-level patch appearance bias while retaining the local structure needed for predictive state learning.

CALVIN ABC-D success rates under different world-model prediction targets.
Prediction target	1	2	3	4	5	Avg. Len.
MAE latent	81.0%	59.8%	44.4%	31.5%	22.6%	2.39
Raw RGB	95.1%	86.7%	78.7%	71.2%	64.3%	3.96
Zero-mean patch	95.1%	88.4%	81.1%	73.2%	66.0%	4.04

CALVIN trajectory variability.

Across 34 CALVIN tasks, per-task mean episode length ranges from 44.6 to 65.0 frames, while within-task standard deviation reaches 12.4 frames. The time-to-go target is normalized independently within each successful trajectory rather than against a fixed global duration.

Instruction Annotation Protocol

This protocol supplies the Goal, Constraint, Action, Spatial, Logic, and Other labels used by the instruction-component masking ablation shown in Results. Each deduplicated CALVIN instruction is labeled automatically, loaded into a span-colored viewer, and checked or corrected by a human annotator. During the ablation, only the corresponding spans in the LGSR encoder’s instruction are masked; the policy always receives the full instruction, isolating representation-level effects.

G Goal
C Constraint
A Action
S Spatial
L Logic
O Other

Read the full annotation prompt

[SYSTEM]
You are an expert annotator for robot-instruction semantics.
You must label tokens/spans into exactly five classes:
- G (Goal): target objects or entities (the noun head of the target NP).
- C (Constraint): attributes, colors, quantities, determiners, ordinal words that restrict G or A.
- A (Action): executable actions (verbs and their particles), e.g., "pick up", "take off".
- S (Spatial): spatial/positional/relational expressions (prepositions, directional adverbs, PP chunks).
- L (Logic): sequencing/logic connectors ("and", "then", "if", "before", "after", ...).
Tokens that fall into none of the five classes are tagged O.

Rules:
1) Prefer span-level labeling with contiguous spans. Multiword expressions like "pick up", "in front of" must be a single span.
2) Within a noun phrase, the head noun is G; its modifiers (color/quantity/ordinal) are C.
3) Spatial phrases that modify G or motion are S (e.g., "in the drawer", "towards the door").
4) Logic words are L and never overlap with A/G/C/S.
5) If a pronoun clearly refers to the current G (e.g. "pick it up"), label the pronoun "it" as G.
6) Output both SPANS and BIO:
   - SPANS: a list of {start, end, label, text} where start/end are token indices (end exclusive).
   - TOKENS: the tokenized sequence.
   - BIO: BIO tags aligned with TOKENS (B-*, I-*, or O); BIO must reflect the same spans.

Tokenization: split by whitespace and punctuation, but keep underscores (e.g., "pick_up"). Do not merge spans across non-adjacent tokens.
Return JSON with keys "tokens", "spans", "bio". No commentary.

[FEW-SHOT EXAMPLE 1]
INPUT: "go towards the pink block in the drawer and pick it up"
{
  "tokens": ["go","towards","the","pink","block","in","the","drawer","and","pick","it","up"],
  "spans": [
    {"start":0,"end":1,"label":"A","text":"go"},
    {"start":3,"end":4,"label":"C","text":"pink"},
    {"start":4,"end":5,"label":"G","text":"block"},
    {"start":1,"end":8,"label":"S","text":"towards the pink block in the drawer"},
    {"start":8,"end":9,"label":"L","text":"and"},
    {"start":9,"end":12,"label":"A","text":"pick it up"},
    {"start":10,"end":11,"label":"G","text":"it"}
  ],
  "bio": ["B-A","B-S","I-S","B-C","B-G","B-S","I-S","I-S","B-L","B-A","B-G","I-A"]
}

[FEW-SHOT EXAMPLE 2]
INPUT: "go to the tower of blocks and take off the top one"
{
  "tokens": ["go","to","the","tower","of","blocks","and","take","off","the","top","one"],
  "spans": [
    {"start":0,"end":1,"label":"A","text":"go"},
    {"start":1,"end":6,"label":"S","text":"to the tower of blocks"},
    {"start":6,"end":7,"label":"L","text":"and"},
    {"start":7,"end":9,"label":"A","text":"take off"},
    {"start":10,"end":11,"label":"C","text":"top"},
    {"start":9,"end":12,"label":"G","text":"the top one"}
  ],
  "bio": ["B-A","B-S","I-S","I-S","I-S","I-S","B-L","B-A","I-A","B-G","B-C","I-G"]
}

[FEW-SHOT EXAMPLE 3]
INPUT: "take off the block that is on top of the other one"
{
  "tokens": ["take","off","the","block","that","is","on","top","of","the","other","one"],
  "spans": [
    {"start":0,"end":2,"label":"A","text":"take off"},
    {"start":3,"end":4,"label":"G","text":"block"},
    {"start":6,"end":9,"label":"S","text":"on top of"},
    {"start":9,"end":12,"label":"C","text":"the other one"}
  ],
  "bio": ["B-A","I-A","O","B-G","O","O","B-S","I-S","I-S","B-C","I-C","I-C"]
}

[INFERENCE]
Now annotate the following sentence:
INPUT: "{sent}"
Return JSON only.

Cross-Benchmark Evaluation

LIBERO-Long contains ten multi-stage tasks lasting roughly 200–400 environment steps. LGSR was trained from scratch on LIBERO-Long demonstrations, using the same architecture and objectives as on CALVIN, without LIBERO-90 or external robot data.

CALVIN ABC-D · shared visual scene

CALVIN environment D scene for closing the drawer. — “push the drawer”

CALVIN environment D scene for grasping the blue block in the drawer. — “grasp the blue block in the drawer”

CALVIN environment D scene for turning off the yellow light. — “turn off the yellow light”

Thirty-four atomic tasks share four environments. Language must disambiguate which skill and object are active within similar visual contexts.

LIBERO-Long · task-distinct scenes

LIBERO kitchen scene for turning on the stove and placing a moka pot. — Stove + moka pot

LIBERO living room scene for placing soup and a box in a basket. — Soup + box + basket

LIBERO study scene for placing a book in the back compartment. — Book + caddy

Longer multi-stage instructions test temporal sequencing, but task-specific scenes make task identity nearly recoverable from vision alone.

LIBERO-Long success rates (%), averaged over the top three checkpoints.
Method	Pretrain	Avg.	T1	T2	T3	T4	T5	T6	T7	T8	T9	T10
MTACT	—	41.0	30.0	50.0	75.0	85.0	20.0	75.0	0.0	0.0	10.0	65.0
MVP	Vision	68.2	83.3	90.0	80.0	88.3	46.7	63.3	45.0	78.3	60.0	46.7
MPI	Vision	77.3	66.6	86.6	96.6	95.0	83.3	83.3	56.6	86.6	40.0	78.3
OpenVLA	VLM	53.7	35.0	95.0	65.0	45.0	40.0	80.0	60.0	45.0	20.0	55.0
Seer (scratch)	—	78.7	80.0	90.0	91.7	81.7	85.0	65.0	86.7	88.3	51.7	66.7
Seer	LIBERO-90	87.7	91.7	90.0	98.3	100.0	91.7	93.3	85.0	88.3	61.7	71.7
LGSR	—	82.3	100.0	91.7	91.7	85.0	83.3	93.3	71.7	90.0	46.7	70.0
LGSR + proxy	—	82.0	88.3	96.7	96.7	93.3	88.3	81.7	58.3	85.0	61.7	70.0

What transfers—and what saturates?

Encoder-stage language fusion transfers: scratch LGSR reaches 82.3%, 3.6 points above scratch Seer. The proxy loss does not improve LIBERO-Long (82.0% versus 82.3%) because the representation probe is already saturated at 100%. On CALVIN, where scenes are shared across tasks, the same auxiliary adds 0.18 Avg. Len. This is why its contribution is described as benchmark-conditional.

Attention Visualization Experiments

Feature-basis probe

A PCA probe first compares the spatial structure of the candidate feature bases. MAE’s pixel-reconstruction features are distributed across the image, while DINOv2 separates objects and coherent scene regions more clearly.

Input CALVIN scene used for feature-basis comparison. — Input

PCA visualization of MAE patch features, which appears spatially distributed. — MAE patch PCA

PCA visualization of DINOv2 patch features, which separates coherent spatial regions. — DINOv2 patch PCA

Why DINOv2 for visualization?

Because its feature regions are spatially coherent, DINOv2 provides a more legible basis for localizing the response of the learned representation. This choice is limited to the visualization probe; LGSR training and quantitative evaluation continue to use the MAE encoder.

DINOv2 attention visualization

The overlays compare the same DINOv2 probe across the raw scene, LGSR without language guidance, and full LGSR, revealing where task conditioning changes the probe response.

Raw CALVIN scene for pushing the pink block right. — Push the pink block right

DINOv2 attention overlay for the model without language on the pink-block task. — Push the pink block right

Raw CALVIN scene for turning off the LED. — Turn off the LED

DINOv2 attention overlay for the model without language on the LED task. — Turn off the LED

Raw CALVIN scene for moving the slider left. — Move the slider left

DINOv2 attention overlay for the model without language on the slider task. — Move the slider left

Raw CALVIN scene for rotating the blue block right. — Rotate the blue block right

DINOv2 attention overlay for the model without language on the blue-block task. — Rotate the blue block right

Scope and Limitations

Successful demonstrations

The reward head is trained only on successful expert clips. Failure rollouts are out of distribution, so the auxiliary is not a calibrated reward for arbitrary policy states.

Alignment requirement

Time-to-go needs a language-aligned clip endpoint. Open-ended or unsegmented tasks would require task-specific milestones such as contact, half-open state, or release.

Benchmark dependence

The proxy objective helps on CALVIN but not on structurally saturated LIBERO-Long. Language fusion is the transferable contribution; the auxiliary benefit depends on task and scene structure.

BibTeX

The accepted preprint does not yet provide volume, issue, pages, or a DOI, so the entry intentionally avoids inventing them.

@article{zhang2026lgsr,
  title   = {Improving Policy Learning via Language-Guided State Representation in World Models},
  author  = {Zhang, Lixuan and Kan, Meina and Shan, Shiguang and Chen, Xilin},
  journal = {IEEE Robotics and Automation Letters},
  year    = {2026},
  note    = {Accepted preprint}
}