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

At a glance

LGSR uses the task instruction to decide which visual information enters a world model’s compact state. A language–vision perceiver and an auxiliary time-to-go prediction loss encourage task-relevant features. On CALVIN ABC-D, the model averages 4.22 completed tasks and 71.9% success across all five tasks, using only the benchmark’s language-annotated demonstrations for robot training.

Why guide the state representation?

A task-agnostic encoder can waste limited capacity on irrelevant objects. LGSR injects language before state compression so the representation emphasizes objects and relationships needed for the instruction. This differs from adding language only at the policy or future-prediction stage. Related work spans world models [2], language-conditioned imitation [24], task-relevant representations [35], and the Seer predictive policy [14].

Problem and learning setup

The agent receives a fixed instruction and multi-view observations with proprioception. It learns from successful expert demonstrations to maximize task completion, using behavior cloning because the binary success signal is sparse and non-differentiable. A shared transformer learns future-state and action predictions; the world model regularizes policy representations rather than serving as a deployed simulator.

Language-guided state representation

Frozen CLIP and MAE encoders produce language and image features. A multimodal perceiver encoder fuses them through learnable latent queries and cross-attention, producing six compact state tokens. The language conditions which scene features survive the bottleneck. The resampler follows the multimodal fusion approach in Flamingo [42].

Figure 1. Framework: frozen MAE and CLIP features enter the multimodal perceiver encoder. The resulting language-guided state supports joint future-observation and action prediction, with an auxiliary reward head.

Equations 1–3

(1)

(2)

(3)

Progress supervision

An auxiliary head predicts normalized time-to-go, decreasing from one at the start of a successful demonstration to zero at completion. Mean squared error trains this head jointly with the policy. The target supplies dense task-related supervision without separate reward annotations, but depends on aligned successful trajectories.

Equations 4–6

(4)

(5)

(6)

Joint world model and policy

The transformer receives ten observations of history, language, the current observation, and learnable slots for future observations and actions. Future slots attend to context; action slots also attend to latent future representations. A shared forward pass produces both. The vision decoder predicts normalized future images, while the action head predicts continuous displacement and discrete gripper state.

Equations 7–10

(7)

e_text = CLIP(l), e_vis = MPE(o_t_vis), e_prop = MLP(o_t_prop)

(8)

(9)

e_f = WM(e_his, e_text, e_vis, e_prop, m_o)

(10)

e_a = PL(e_his, e_text, e_vis, e_prop, e_f, m_a)

Training objectives

Future reconstruction uses patch-wise zero-mean targets. Continuous actions use mean squared error; gripper commands use cross-entropy. The final objective combines action, future-observation, and proxy-reward losses with weights 1.0, 0.1, and 0.01; the discrete gripper term has weight 0.01. Future prediction and action chunks span three steps.

Equations 11–14

(11)

o_hat_(t+N) = VD(e_f), a_hat_(t:t+M) = AH(e_a)

(12)

(13)

L_a = MSE(a_hat_cont, a_cont) + w_disc CE(a_hat_disc, a_disc)

(14)

Evaluation and implementation

Training uses CALVIN environments A–C and evaluation uses unseen environment D [15]. One thousand instruction chains each contain five subtasks. Reported results average the top three checkpoints selected by validation performance. The model contains 320 million parameters, with 69 million trainable. Training takes approximately 192 GPU-hours on RTX 4090 hardware; reported inference is about 40 Hz. Frozen-feature task identification is much harder on CALVIN than LIBERO-Long, motivating this benchmark choice.

Table I. Task identification accuracy (mean ± standard deviation) on CALVIN ABC-D and LIBERO-Long for Seer, LGSR, and LGSR with proxy reward.

Main results

LGSR with proxy reward reaches 4.22 average completed tasks and 71.9% five-task success, compared with 3.64 for Seer trained from scratch and 3.98 for pretrained Seer. Language guidance alone reaches 4.04. LGSR exceeds GR-MG on multi-step completion while using fewer trainable parameters and no robot-data pretraining. These comparisons involve different architectures and pretraining data; Table II preserves those distinctions.

Table II. CALVIN ABC-D success for one to five consecutive instructions, average length, trainable parameters, and training-data sources. Results use 1,000 instruction chains. D_l and D_nl denote language-annotated and unannotated CALVIN data; D_e denotes Ego4D and D_OXE denotes Open X-Embodiment.

Data, capacity, and instruction ablations

At 70% of the language-annotated data, LGSR approaches its full-data performance and exceeds Seer trained with full annotated and unannotated data. Seer retains an advantage at 20% data or below. Language guidance helps across one to six representation tokens; with one token, the relative gain is 17.7%. Masking goal, constraint, or spatial words hurts performance most; removing some logical or incidental words can help.

Figure 2. Ablations of training-data proportion, representation-token count, and instruction components. Goal, constraint, spatial, action, logic, and other words are masked separately; error bars show the standard error of the mean.

Which progress target works?

The time-to-go target reaches 4.22 average task length, compared with 3.93 for elapsed progress and 3.76 for a binary terminal-success target. The authors attribute this to stronger task-discriminating information in early states. This is an auxiliary representation-learning result, not a tested reward for failure recovery.

Table III. Proxy-reward ablation: binary success, elapsed progress, and time-to-go targets, evaluated over one to five CALVIN instructions and average sequence length.

What information remains?

Reconstructions from a frozen one-token state retain task-relevant blocks or the drawer while discarding appearance details. A separately trained inverse decoder maps this state back into MAE features before the frozen image decoder reconstructs it. The visualizations indicate selective information retention; they are not direct measurements of attention or causal reasoning.

Figure 3. Raw images, MAE reconstructions, and one-token MPE reconstructions for lifting blue, pink, and red blocks and closing a drawer.

Scope and limitations

The reward head only sees successful demonstrations, so behavior on failures is underexplored. Time-to-go needs alignment between task language and trajectories or milestones, and training still requires costly language-annotated expert data. Benefits depend on the benchmark: the auxiliary objective gives near-zero improvement on LIBERO-Long. The paper proposes failure-aware supervision and less annotated data as future work.