Improving robot planning by using time-aware world models
Exploring Temporal Context-Aware Physical Grounding in Latent World Models
Artificial Intelligence
Summary
Robots need to understand how their actions will change the world to plan well. The authors improve existing robot learning methods by teaching the system to better focus on important physical details over time. They also help the robot choose smarter actions by considering how its current state compares to the goal. This makes the robot better at planning tasks, especially when there is a long wait between the start and goal moments.
What this means in practice
- •For robotics engineers: Create robots that plan complex manipulation tasks by learning time-aware physical models for better action choices.
- •For robotic software developers: Build robot controllers that use improved latent models to plan actions more effectively over long task timelines.
Authors
Yilei Man, Yang Chen, Jingcai Guo
Abstract
Learning latent dynamics is fundamental to building action-conditioned world models for planning. Among existing families, Joint Embedding Predictive Architectures (JEPAs) offer a natural framework by modeling transitions directly in latent space rather than at the pixel level. However, two issues remain in JEPA-based world models: (i) physical grounding lacks explicit aggregation of temporal context and (ii) action planning always overly relies on generic sampling distributions. To this end, we propose $\mathcal{G}$-JEPA, which incorporates temporal context into physical grounding for latent world models and couples the learned models with prior-guided action planning. Specifically, during learning, we incorporate a shared causal temporal module that aggregates temporal context within observation and prediction sequences under physical state supervision, grounding both in task-relevant physical states. When planning, we further introduce action inference through stepwise conditioning of a generative action distribution on the current and goal latent representations together with their evolving discrepancy, yielding anchored proposals that guide model-based online refinement. Together, physical state supervision guides temporal representation learning, while the action prior guides search through the learned world model. Experiments demonstrate improved planning performance on manipulation tasks, including tasks with increasing temporal separation between initial and goal observations. Code is publicly available at https://github.com/myl11045/G-JEPA.