Planning is terribly slow

World models has recently become the highly overloaded buzz word for anything that leads to real world intelligence; from video generation models to predictive architectures that use MBRL or MPC, and everything in between. One of these methods, JEPA-WM, plans the optimal trajectory in imagination before acting. They can perform actions very fast (in fact, as fast as hardware can support) until alas! We have veered far off from the imagination, and need to do the expensive planning again, from this state on. For perspective, DINO-WM, with ViT-S predictor and MPC-CEM planner, thinks for 2mins on an RTX PRO 4500 for the push-t task, which is then performed in 2 seconds. VLAs on the other hand, use a greedier approach, thinking hard at every step, making their actions itself very slow. For perspective, folding a shirt takes several minutes with current VLAs.

We explore 3 ways to make planning faster:

We can cache redundant information

From classic algorithms theory, we know that there is always a tradeoff between time and space complexity. Thus, one way to make planning fast is to make use of space. Current methods, especially those based on monte-carlo sampling of rollouts (CEM), need to pass the same initial state and then slightly different subsequent states through the predictor for each rollout. This process can be sped up if we "cache" certain patches which don't vary much from one state to the next, allowing us to essentially focus or gaze (Autogaze) on only the important part of the current state which is most likely to change after taking an action. Taking inspiration from the popular KV caching techniques being used to speed up transformers in LLMs, we could cache the projection of these patches in the query, key and value spaces. In fact, we could even learn more composable object-level representations (C-JEPA), so that we can cache objects, instead of patches.

Make the search for optimal trajectories faster

There has been recent work (GRASP) on parallelizing planners by jointly optimizing over virtual intermediate states and actions, using latent collocation and myriad of other tricks to stabilize training. However, we feel a graver problem is that we start from random samples, hence make convergence slow. If we can use appropriate priors that are greedily found using only the initial and goal state, followed by optimization, the planning will be much faster. We could also use combinations of existing methods, such as CEM-GD. The former proposal relies heavily on the fact that state space preserves notion of distance, which may not be the case. Another method that may work is to train a model (such as diffusion), that gives a high level, inaccurate estimate of the trajectory, in one step as opposed to doing it auto-regressively, which is then used as a prior.

The state space is conducive for planning

The joint embedding architecture gives a state space that is generally good for downstream tasks. We need to make it more conducive for our specific task, i.e., planning. Current methods work on continuous latent spaces, which lead to noisy trajectories, that some approaches are now trying to solve by ensuring local linearity (temporal straightening). Good representations on higher dimensions are sparse (Rectified LpJEPA), allowing for different concepts to be separated nicely. A method that enforces this sparsity could allow us to then discretize the latent space (Dreamer v4). This discretization then has the effect of making the space locally linear and thus allowing nice, smooth trajectories to be sampled. Such a state space, along with a discrete action space will also allow more human-like planning, which is more focused that just random searching. Using dynamic programming methods like A* search, one can include heuristics which act like a good prior direction along which to move. We claim that all these are means to the same task.

Currently, we are focussing more on the caching method, as it has immediate gains in the raw clock time, and has been shown to work in NLP. However, the largest gains will come from planning in a more optimal search space. Long term memory, trainable cost, sim2real, configurator, in our opinion, should only be dealt when planning is no longer a bottleneck.