World Models @ IITB

Embodied AI Research

Breadth-First Search

24 March 2026

Papers/blogs/pods consumed

CEM-GD

Uses CEM trajectory as an initialization for Gradient Descent; Achieved similar success rate as CEM in half the time.

GRASP

Optimizes over "virtual" states and actions to find the optimal trajectory in goal conditioned tasks, achieving exceptional parallelism.

Dreamer v4

A mashed up set of tricks from NLP, CV and SSL that were arduously adapted for mining diamond in Minecraft, using ~90k hrs of video of people playing Minecraft.

Deep Hierarchical planning from Pixels

Manager policy (system 2) outputs sub-goals in a discrete code space in order to maximise goal and exploration rewards, and the worker policy (system 1) tries to reach the sub-goal.

notboring blog

A brief history of world models: from the Matrix movies to SOTA embodied agents like SIMA 2 and $\pi0.6$.

Saining Xie pod

7 hours marathon on pure world model research and splendor.

Why AI systems don't learn and what to do about it

Learning must be a core capability, rather than an architectural choice outsourced to human experts.

Temporal Straightening

A cognitive idea from Neuroscience makes its way to WMs, enabling "nicer" state spaces for planning.

Le-WM

Swap in a SIGREG regularizer into classic JEPA-WMs, and train Encoder + Predictor jointly, achieving comparable performance to SOTA with lesser model complexity, and nicer state spaces as a by-product.

V-JEPA 2.1

Learn dense features in videos and images (a first!) that results in very satisfying segmentation and object tracking results (a lot of them!).

Ideas discussed

Uncertainty loss in planners

Inspired by Uncertainty Aware RWM. Have the predictor output mean and variance of the next state (similar to VAEs), and add a variance term to the MPC objective. Uncertainty prediction assists to find safe trajectories, can it find them faster than current methods?

Sampling based planning like RRT

Tree based search methods for efficient exploration of high dimensional state space, with in-built memory and forgetting mechanisms. How to metrize state space for making such planners viable? Will it work in a non-markovian setting?

Memory graph

Storing long term memory as nodes with state, action and task tokens. How to retrieve memories? How to keep size tractable?

MTP and Speculative decoding inspired method

Use a draft model to predict multiple states into the future, and use MPC for waypoint navigation through those states. How to keep the draft model small? How to build this mathematically, beyond just an architectural tweak?

Spatial Caching

Only changed patches can query, while Key, Value pairs of redundant patches are cached. How to decide redundant patches with a smaller learned model like Autogaze? How to make it work in non-static camera setups?

Predictive Coding

The human brain certainly does not do backprop. Predictive coding is an interesting alternative that agrees with many neuroscience principles like continual learning, active inference, hebbian plasticity, etc.

Distance preserving state space by following geometric structure of the action space

An isometric loss term that causes the distance between two states to be dependent on how many/what kind of actions one needs to take to get from one to other. How do we find a provable loss metric that enforces this without collapse, and considers the insanely compressed action space? Condition it on state history or task tokens?

Tasks completed

CEM-GD implementation

We implemented CEM-GD and found it to be twice as fast as CEM, with comparable success rate. See here.

GRASP implementation

We implemented GRASP, but couldn’t replicate scores claimed in the paper despite extensive hyper-parameter sweeps. See here.

Dreamer v4 implementation

We began implementing Dreamer v4 in simple environments like Two Rooms and PushT, on top of the Stable World Models repo.

To conclude…

We delved deeper into the ideas previously discussed, implemented some of them, and formulated the rest more rigorously. A lot of questions still remain, and we hope to find answers, through theory and experiments in the upcoming weeks. We are also cooking up a cool demo, that we hope demonstrate our proposed novelties on.