FLARE: Robot Learning with Implicit World Modeling

Ruijie Zheng; Jing Wang; Scott Reed; Johan Bjorck; Yu Fang; Fengyuan Hu; Joel Jang; Kaushil Kundalia; Zongyu Lin; Loic Magne; Avnish Narayan; You Liang Tan; Guanzhi Wang; Qi Wang; Jiannan Xiang; Yinzhen Xu; Seonghyeon Ye; Jan Kautz; Furong Huang; Yuke Zhu; Linxi Fan

Source-linked AI summary

FLARE: Robot Learning with Implicit World Modeling

Ruijie Zheng, Jing Wang, Scott Reed, Johan Bjorck, Yu Fang, Fengyuan Hu, Joel Jang, Kaushil Kundalia, Zongyu Lin, Loic Magne, Avnish Narayan, You Liang Tan, Guanzhi Wang, Qi Wang, Jiannan Xiang, Yinzhen Xu, Seonghyeon Ye, Jan Kautz, Furong Huang, Yuke Zhu, Linxi Fan

arXiv:2505.15659v1May 21, 2025cs.ROcs.LG

TL;DR

Robot policies often model actions alongside future visual frames, but pixel-level prediction is computationally costly and can compete with compact action representations. FLARE instead aligns diffusion-transformer features with future observation embeddings using a lightweight latent world-modeling extension, achieving state-of-the-art manipulation results and enabling action-free video co-training.

Problem
Joint future-frame and action learning requires costly high-fidelity visual prediction and combines competing demands for spatially detailed versus compact task-relevant representations.
Method
FLARE adds learnable future tokens to a diffusion transformer and aligns their hidden representations with future observation embeddings while retaining action flow matching.
Results
FLARE achieves state-of-the-art performance across two multitask manipulation benchmarks and supports video-only human demonstrations, with 95% success after 100 real GR1 trajectories per task.
Takeaways & Limitations
FLARE provides a lightweight way to combine implicit future-state modeling with robot policy learning and co-training from action-free human egocentric videos.
Takeaways & Limitations
The study focuses mainly on imitation learning for pick-and-place tasks on a real humanoid robot and still relies on a small number of expert demonstrations.

Abstract

from arXiv · show

We introduce $\textbf{F}$uture $\textbf{LA}$tent $\textbf{RE}$presentation Alignment ($\textbf{FLARE}$), a novel framework that integrates predictive latent world modeling into robot policy learning. By aligning features from a diffusion transformer with latent embeddings of future observations, $\textbf{FLARE}$ enables a diffusion transformer policy to anticipate latent representations of future observations, allowing it to reason about long-term consequences while generating actions. Remarkably lightweight, $\textbf{FLARE}$ requires only minimal architectural modifications -- adding a few tokens to standard vision-language-action (VLA) models -- yet delivers substantial performance gains. Across two challenging multitask simulation imitation learning benchmarks spanning single-arm and humanoid tabletop manipulation, $\textbf{FLARE}$ achieves state-of-the-art performance, outperforming prior policy learning baselines by up to 26%. Moreover, $\textbf{FLARE}$ unlocks the ability to co-train with human egocentric video demonstrations without action labels, significantly boosting policy generalization to a novel object with unseen geometry with as few as a single robot demonstration. Our results establish $\textbf{FLARE}$ as a general and scalable approach for combining implicit world modeling with high-frequency robotic control.

1 Introduction

FLARE adds latent future-state prediction to diffusion or flow-matching robot policies without full-frame reconstruction. It improves multitask manipulation performance and supports training with action-free human egocentric videos.

Motivation: FLARE addresses the computational overhead and competing representation demands of jointly predicting future visual frames and actions.High-fidelity visual prediction requires large generative models, while action modeling benefits from compact, task-relevant representations.
Approach: FLARE introduces a future alignment objective that predicts compact future-observation representations from the action-denoising network.The method avoids full-frame reconstruction and adds only a lightweight extension to diffusion or flow-matching policies.
Approach: FLARE combines action flow matching with future latent alignment, unlike a conventional flow-matching or diffusion policy.This joint training enables the policy to learn from both action-labeled demonstrations and video-only human egocentric demonstrations.
Results: FLARE achieves state-of-the-art performance across two multitask imitation-learning benchmarks covering single-arm and humanoid tabletop manipulation.The reported evaluations also include real humanoid learning and generalization from human egocentric video with minimal robot data.
Results: With one real robot demonstration per object, FLARE learns novel grasping strategies from GoPro-collected human egocentric videos without action labels.The demonstrations support generalization to novel objects with unseen geometry.

2 Background

The policy uses flow matching to generate action chunks conditioned on current vision-language and proprioceptive observations. Its DiT architecture iteratively denoises actions through a small number of steps.

Flow-matching policy: Flow matching fits action chunks from human demonstrations by conditioning on the robot’s vision-language embedding and proprioceptive state.The observation includes image inputs and a language instruction, while the action chunk contains a horizon of future actions.
FLARE context: FLARE extends this policy architecture by processing state, action, and future-token embeddings together in the DiT sequence.The future-token activations are later aligned with future observation representations.
Flow-matching policy: The policy generates action chunks by sampling an initial noisy chunk and refining it through forward Euler integration.Inference uses iterative denoising rather than direct one-step action prediction.
Flow-matching policy: The DiT alternates cross-attention and self-attention layers while conditioning on the robot’s vision-language embedding.The implementation uses K = 4 denoising steps throughout the experiments.

3 Method

FLARE adds learnable future tokens whose intermediate DiT representations are aligned with future observation embeddings while action flow matching remains active. An action-aware compact embedding model supplies efficient targets for this alignment.

Future-token architecture: FLARE concatenates encoded proprioceptive state, noised action chunks, and M learnable future tokens before processing them with DiT blocks.The future tokens provide a separate stream for predicting future latent states.
Future-token architecture: Intermediate future-token activations are projected with an MLP and aligned with frozen vision-language embeddings of the future observation.The activations are extracted at an internal layer L, producing the inputs to the future latent alignment objective.
Design distinction: FLARE differs from REPA by aligning policy representations with future rather than current embeddings and by adding separate interacting future-token and flow-matching streams.The streams interact through self-attention while retaining action prediction.
Training objective: The combined objective preserves action prediction through flow matching while encouraging internal reasoning about future latent states.The alignment objective uses the DiT future-token activations and future-observation encoder outputs.
Training objective: λ = 0.2 produced the best empirical performance in the reported experiments.The coefficient weights the future latent alignment term in the overall loss.
Action-aware embeddings: The action-aware embedding model uses SigLIP-2 vision and text encoders, multimodal self-attention, and a Q-former to produce M = 32 query tokens.This yields a compact fixed-size representation that supports multi-camera inputs and is optimized for policy learning.

4 Experiments

The experiments evaluate FLARE across multitask simulation, real humanoid manipulation, cross-embodiment post-training, human-video co-training, and training-stability ablations. FLARE consistently improves policy performance, including under limited robot data and with action-free human demonstrations.

Multitask Benchmark Performance: 24 RoboCasa and 24 GR-1 tabletop tasks form the two multitask simulation benchmarks, covering single-arm and humanoid manipulation.RoboCasa includes kitchen tasks, while GR-1 emphasizes dexterous humanoid control.
Multitask Benchmark Performance: FLARE consistently outperforms policy-only and video-action prediction baselines across the multitask benchmarks.UWM receives five times the training budget of the other methods, while FLARE’s policy-only variant does not match the full objective.
Cross-embodiment Post-training: 10% gain on RoboCasa with 100 trajectories per task demonstrates stronger performance under limited post-training data.The pretrained embedding also reaches 71.3% versus 70.2% for an in-domain embedding with 1000 trajectories.
Cross-embodiment Post-training: 95.1% success rate is achieved on the real GR-1 humanoid, averaging 14% higher than the baseline method.FLARE maneuvers around or over nearby objects that the policy-only baseline often knocks over.
Human Egocentric Trajectories: 80% success rate on novel objects is reached with 10 trajectories per object when FLARE is jointly trained with human egocentric videos.With only 1 teleoperated trajectory per object, FLARE reaches up to 60% success.
Ablation Study: ρ = 0.995 gives the best EMA performance, while every tested EMA variant outperforms the baseline without future latent alignment.ρ = 1.0 still surpasses the baseline, whereas ρ = 0.99 performs worst, likely because of frequent target updates.

5 Related Work

FLARE extends latent representation alignment to robot policy learning, distinguishing itself from explicit visual prediction, latent-action intermediates, and pose-based video methods. It uses future, action-aware observation embeddings to support compact policy and world-model co-training.

FLARE’s distinction: FLARE aligns policy representations with future observations rather than current observations, introducing an implicit latent world model without explicit frame reconstruction.The approach adds learnable future tokens and separates flow matching from alignment streams that interact through self-attention.
Scope: FLARE is designed for joint policy and world-model co-training, in contrast to DINO-WM’s focus on zero-shot planning.Planning is identified as a possible future extension rather than the framework’s current focus.
Human video learning: Unlike pose-estimation and motion-capture approaches, FLARE learns from action-free egocentric videos without explicit pose estimators or point tracking tools.Its future latent alignment is presented as a lightweight alternative for extracting temporally useful information.
Representation choice: FLARE directly aligns action-aware vision-language embeddings with future observations, whereas latent-action methods use intermediate representations whose relation to ground-truth actions is unclear.The direct alignment is described as capturing temporal dynamics relevant to policy learning.

6 Limitations

The paper’s evaluation centers on imitation learning for pick-and-place tasks on a real humanoid robot, with broader task, data, and learning settings left for future work.

Current scope: The real-robot study focuses mainly on imitation learning with pick-and-place tasks on a humanoid robot.More fine-grained dexterous manipulation and reinforcement-learning integration remain future directions.
Data requirements: FLARE still relies on a small number of expert demonstrations, which may limit scalability where such data is difficult to acquire.This boundary remains despite the method’s ability to generalize to novel objects.
Video-data scope: The egocentric-video experiments use controlled datasets collected with head-mounted GoPro cameras rather than diverse natural-environment motion data.Larger and more varied datasets in natural environments are proposed as future work.

7 Conclusion

FLARE jointly learns robot policies and latent world dynamics by aligning future observation representations with action-denoising states. The resulting framework reaches strong manipulation performance and supports action-free human-video co-training for novel-object generalization.

Framework: FLARE jointly learns robot policy and latent world dynamics through future representation alignment.The framework aligns future observation representations with hidden states of the action denoising network.
Framework: FLARE enables policies to implicitly reason about future states while predicting actions.This capability is obtained without requiring explicit future-frame reconstruction.
Results: FLARE achieves state-of-the-art performance on challenging robotic manipulation benchmarks.The conclusion presents this as an outcome of the future latent alignment framework.
Results: FLARE supports co-training with action-free human egocentric videos and improves generalization to novel objects with minimal real-robot teleoperation data.The conclusion links the video capability to novel-object generalization.

B Pretraining Data Mixture

This section identifies the pretraining data mixture through a dataset-statistics table, but the supplied passages provide no mixture contents or numerical statistics.

Dataset statistics: The pretraining data mixture is presented in Table 3.No dataset components or quantities are specified in the supplied passage.
Dataset statistics: Table 3 is titled “Action-Aware Vision Language Embedding Pre-training Dataset Statistics.”The supplied text references the table but does not include its rows or values.

C Training Details

FLARE training combines action flow-matching with future-representation alignment in a diffusion-transformer policy. The implementation adds future tokens and optimizes both objectives in a repeated training loop.

Input preparation: Each iteration creates noisy actions, samples a flow-matching timestep, and computes the action velocity target.The noisy action interpolates between the action and Gaussian noise.
Policy computation: The policy receives state, action, and observation embedding tokens together with future tokens through the diffusion-transformer layers.State and action tokens are concatenated with future tokens, while vision-language tokens condition the transformer.
Action objective: The action objective minimizes mean-squared error between decoded actions and the target velocity.Action outputs are decoded from the policy outputs corresponding to the action tokens.
Future alignment: The FLARE objective minimizes one minus cosine similarity between predicted embeddings and target embeddings of future observations.The target future-observation embedding is computed without gradient tracking.
Optimization: Training optimizes the combined loss = action_loss + lambda * flare_loss using gradient backpropagation and an optimizer step.The pseudocode repeats optimization for N gradient steps.

E.1 4 Pick-and-place Tasks

The paper presents FLARE rollouts on a real GR1 humanoid robot across four pick-and-place tasks. The qualitative examples show object-aware hand maneuvering during manipulation.

Qualitative behavior: When manipulating a bottled water or Coke can, FLARE maneuvers the hand around the object rather than striking it.For the water bottle, the hand hovers over the object instead of knocking it over.
4 Pick-and-place Tasks: Four real-world GR1 humanoid pick-and-place tasks are included in the FLARE policy rollout demonstrations.The figure is presented together with each task’s language instruction.

E.2 Manipulating Novel Objects

The paper presents FLARE policy rollouts manipulating five novel objects, demonstrating evaluation beyond the objects used in the primary demonstrations.

Manipulating Novel Objects: FLARE policy rollouts evaluate manipulation of five novel objects.The passage identifies these as novel-object policy rollouts without reporting success values.