At Tau, we are building a general AI for robots. This requires a massive amount of data for training. Learning from human teleoperation demonstrations is not sufficient. Robots will have to learn from data without labeled actions through reinforcement learning on trillions of trials.

In the past, humans have developed physics simulators to generate this data. However, developing new environments is labor-intensive and some aspects like deformable objects or other humans remain hard or impossible to simulate.

Another approach is to learn a simulator on large amounts of unlabeled real world data. This scales well with more compute and data, and is not bound by human labor.

Model Architecture

In this report, we introduce a latent autoregressive flow-matching world model (LARF) and scale it to 1B parameters. It predicts real world images frame by frame in a compressed latent space with a large causal transformer and a small flow-matching head¹¹⁹. This allows us to inference the causal transformer once and use the output to predict the next observation and a future action which significantly speeds up policy optimization in the learned world model.

Frame-Wise Latent Space

To encode the high-dimensional pixel space of images into a low-dimensional continuous latent space, we design a frame-wise variational autoencoder (VAE)³. Given an input video V ∈ R^T×H×W×C, the tokenizer compresses each frame in the spatial dimensions to [H/32, W/32] while expanding the number of channels C to 24.

Some of our datasets contain observations with multiple camera angles. We concatenate the images channel-wise before compressing them. We train the VAE to minimize the L1 reconstruction loss, KL loss, LPIPS perceptual loss⁴⁵as well as GAN loss⁶.

While increasing the channel dimension achieves improved reconstruction performance measured in PSNR, LPIPS and SSIM, it leads to degraded generation quality. To overcome this optimization dilemma, we use semantic regularization of the latents to DinoV2⁷features, by introducing a cosine similarity loss as well as a distance similarity loss⁸⁹¹⁰.

Causal Transformer

We encode the sequence of latents with a causal transformer model. It applies attention alternately over the time and space dimensions. The temporal attention treats width and height as batch dimensions and computes attention only over the temporal dimension after adding a one dimensional rotary position embedding¹¹. It applies causal masking so that key-value caching can be used at inference time. We further use grouped query attention¹²to reduce the size of the key-value caches. The spatial attention is per timestep over the width and height axes and uses a two dimensional rotary position embedding¹³. We use RMSNorm¹⁴and SwiGLU¹⁵for the feedforward layers.

Unfortunately, autoregressive prediction is prone to error accumulation. To mitigate this drift, we add a random amount of noise to the latent frames at training time¹⁶. A learned embedding is used to input whether the latent frame is noisy. At inference time, we use the clean embedding for the real context frames and the noisy embedding for the generated latent frames.

This noise embedding combined with actions if available is inserted into the transformer via AdaLN layers².

Flow-Matching Head

The output of the encoder is normalized and added to the conditioning of the flow-matching head¹⁷¹⁸which is implemented using a standard diffusion transformer architecture with only spatial attention². The objective is to predict the next latent frame. However, we noticed that predicting the delta between consecutive frames performs better than direct prediction.

At training time, we sample a noise timestep t independently for each episode timestep i from a logit-normal distribution¹⁹. We construct the training sample x^t = t (x_i+1 - x_i) + (1-t) x⁰, where x⁰ ∼ 𝒩(0, 1) and x_i correspond to the sequence of latents x_i after adding a random amount of noise to mitigate error accumulation. The model is then trained to minimize the mean squared error to the velocity v^t = (x_i+1 - x_i) - x⁰. At inference time, we use a simple first-order Euler ODE solver with uniformly distributed timesteps.

Experiments

We curated a dataset sourced from both internal sources and publicly accessible data. Our in-house dataset is collected mainly autonomously by our robots and comprises random exploration data as well as task-specific data such as picking, grasping and sweeping. For our internet-sourced dataset, we use the Ego4D dataset²⁰and royalty-free videos from Pexels as well as RoboSet²¹. The first two sources depict human activities whereby Ego4D contains data captured with a head-mounted camera to show observations from an egocentric perspective.

In total, we train on around 0.5B frames. All videos are sampled at 10 fps with a resolution of 256px after filtering out low-quality clips.

We scale our model up to 1B parameters and pre-train on the entire dataset for 42 hours on 64 H100 GPUs. Afterwards, we fine-tune exclusively on robot data for 26 hours on 16 H100 GPUs. Inference speed for autoregressive latent frame prediction on a single H100 is around 110 fps at a batch size of 32. Decoding the latent frames to pixels is only needed for us humans. Policy optimization will happen in the latent space. Below you can see some generations on human and robot data.

Future Work and Limitations

Our 1B parameter model generates accurate predictions for simple scenes, while generating complex collisions or other humans remains challenging. We expect significant performance improvements by expanding model size and dataset scale.

However, simulating complete visual observations is nearly always unnecessary for reinforcement learning. Images contain a lot of task-irrelevant information. Using a task-relevant latent space is an exciting future direction for improved efficiency of world model and RL training.

Another challenge is long-horizon tasks where autoregressive frame-wise simulation is slow and prone to significant error accumulation. To address this inefficiency, future work should investigate hierarchical approaches that enable both world models and RL policies to operate in temporally abstract spaces.

References