google-deepmind/motion-forecasting

Python

Open original ↗

Captured source

source ↗

google-deepmind/motion-forecasting

Language: Python

License: NOASSERTION

Stars: 0

Forks: 0

Open issues: 0

Created: 2026-10-07T10:16:48Z

Pushed: 2026-10-07T13:01:02Z

Default branch: main

Fork: no

Archived: no

README:

Forecasting Motion in the Wild

Neerja Thakkar1,2, Shiry Ginosar3, Jacob Walker2, Jitendra Malik1, Joao Carreira2, Carl Doersch2

1 UC Berkeley 2 Google DeepMind 3 TTIC

ECCV 2026

[[Paper (Arxiv Version)]](https://arxiv.org/abs/2604.01015) [[Project Page]](https://motion-forecasting.github.io/)

Overview

We propose dense point trajectories as visual tokens for behavior, a structured mid-level representation that disentangles motion from appearance and generalizes across diverse non-rigid agents. Building on this abstraction, we design a diffusion transformer that models unordered sets of trajectories and explicitly reasons about occlusion, enabling coherent forecasts of complex motion patterns.

Given a single input image and a short motion history, our model forecasts future animal motion as a set of point trajectories.

Installation

# Clone the repository
git clone https://github.com/google-deepmind/motion-forecasting.git
cd motion-forecasting

# Install dependencies
pip install torch torchvision timm hydra-core omegaconf lightning tqdm scipy

DINOv3 Setup

Our model uses frozen DINOv3 ViT-L features for per-point visual conditioning. You need:

1. A local clone of the DINOv3 repository 2. The pretrained weights (dinov3_vitl16_pretrain_lvd1689m-8aa4cbdd.pth)

Set the environment variable before training or inference:

export DINOV3_REPO_DIR=/path/to/your/dinov3

Inference

Expected workflow: you have a video, cropped it (square) around the animal, and ran a point tracker (e.g. TAPIR) on the crop — giving per-point (x, y) coordinates and visibility over time. Feed the model the first frame of the crop plus the observed track history; it forecasts the future motion.

python scripts/inference.py \
--checkpoint /path/to/checkpoint.ckpt \
--image /path/to/image.png \
--tracks /path/to/tracks.npz \
--output predictions.npy \
--visualize --viz_output viz.png

Input format:

  • --checkpoint: Released checkpoints are self-describing (they embed the model config), so no config file is needed. For other checkpoints, pass --config (the training run's config.yaml) to define the architecture.
  • --image: First frame of the crop (PNG/JPG). Must be a square crop around the animal — it is resized to the model's 256×256 input internally, so non-square images get squashed and distort the predicted motion.
  • --tracks: Bundled .npz with tracks (T, N, 2) and optionally visibility (T, N) (1=visible; defaults to all-visible). Coordinates in [0, 1] normalized to the image (pixel coordinates are auto-detected and normalized). T must be at least the observed history length (num_point_cond, default 4). Any N works: more points than the model's 320 are uniformly downsampled, fewer are padded internally and stripped from the output. A plain (T, N, 2) .npy is also accepted.

Displacement conditioning: unconditional by default. Pass --displacement DX DY to steer the forecast — values are in pixels of the 256x256 model crop (64 0 = move right by a quarter of the image, 256 256 = the full image diagonal).

Demo examples: the released demo examples (one folder per example with image.png, tracks.npz, meta.json) are in exactly this format — point --image/--tracks at any of them. The format is documented in [docs/DATA_PREPROCESSING.md](docs/DATA_PREPROCESSING.md).

Output:

  • predictions.npy: Predicted future tracks, shape (T, N, 2)
  • viz.png (optional): Visualization of observed + predicted tracks overlaid on the image

DDIM options:

  • --use_ddim / --no_ddim: Toggle DDIM sampling (default: on)
  • --ddim_steps 50: Number of DDIM steps
  • --ddim_eta 0.0: DDIM stochasticity parameter

Training

Training uses Hydra for configuration and Lightning Fabric for distributed training (DDP).

The training config is conf/train.yaml; any key can be overridden on the command line (Hydra syntax):

python scripts/train.py data_root=/path/to/examples epochs=10

data_root points to a folder of demo-format examples (one subfolder per example with image.png + tracks.npz; see [docs/DATA_PREPROCESSING.md](docs/DATA_PREPROCESSING.md)).

To finetune from the released checkpoint, set load_path (see the comments in conf/train.yaml).

Training with a custom dataloader

The training loop accepts any dataloader that yields batches as a dict with these keys:

| Key | Shape | Description | |-----|-------|-------------| | video | (B, T, C, H, W) | RGB frames, float [0, 255] (T=1 frame is used) | | tracks | (B, T, N, 2) | Point tracks, normalized [0, 1] to the crop | | visibility | (B, T, N) | 1=visible, 0=occluded (optional) | | point_mask | (B, N) | True=real point, False=padding (optional) | | total_displacement | (B, 2) | Mean per-point displacement (optional) |

To use your own dataloader, replace the ExampleDataset / DataLoader construction in engine/train.py with anything that yields this dict.

Evaluation

Evaluation is configured by conf/eval.yaml; any key can be overridden on the command line (OmegaConf dotlist syntax). It runs over a folder of demo-format examples (data_root); released checkpoints embed their model config, so only the checkpoint path is needed:

python -m engine.eval \
checkpoint=/path/to/model.ckpt \
data_root=/path/to/examples ddim_steps=100

# With paper-style visualization:
python -m engine.eval checkpoint=/path/to/model.ckpt data_root=/path/to/examples \
visualize=true viz_style=animal_color viz_dir=viz_out noise_seed=0

Pretrained Checkpoints

| Model | Dataset | Download | |-------|---------|----------| | DiT-B (Ours) | MammalMotion | Coming soon |

MammalMotion Dataset

We release MammalMotion, a large-scale dataset of camera-stabilized point trajectories...

Excerpt shown — open the source for the full document.