Kubric: A scalable dataset generator

In plain words

Kubric is a tool for making training and test data for vision models out of 3D scenes. Instead of collecting and labelling photos by hand, you write a short script that places random objects on a floor, lets a physics engine move them, and renders the result with a ray tracer. Because the scene is fully known, the tool can also output the true depth, which object each pixel belongs to, and how each pixel moves from one frame to the next.

Each scene is independent, so the same script can run on one laptop or be spread over thousands of machines. The paper uses the tool to build a set of datasets for different problems, from finding objects in video to tracking points through a clip and rebuilding 3D shapes from images. In the authors' tests, optical-flow pretraining on this synthetic data beats the older FlyingChairs set, while a learned-rendering method still does better.

TL;DR

Kubric is an Apache-2.0 Python framework (Google Research) that sits on top of PyBullet for rigid-body physics and Blender's Cycles path tracer §3.1; §3.2. Each worker script builds and renders one random scene, and many workers run in parallel §3.2. The paper's contribution is the tool plus a set of dataset use cases; the numbers it reports are author-run, and several training protocols are not stated §4. Read it as a reference for the worker pattern and the MOVi scene generators, and treat the benchmark numbers as context.

Problem

Dense ground truth, such as optical flow, cannot be obtained on generic real video, even by human annotation §4.2. Real data at scale is costly to annotate and can raise privacy, fairness and licensing concerns §1. Synthetic data avoids these costs and offers full control over content §1. The authors argue that the usual gap to real data is narrowed by photo-realistic rendering §1; §3.1.

Earlier generators were usually one-off, task-specific pipelines, and many general frameworks were rasterization-based or lacked some annotations §1; §2. The authors' stated goal is a common framework for photo-realistic generation with shared assets, reuse and scaling §1.

Costs the paper reports: path tracing can be several orders of magnitude slower than rasterization §3.2, and the authors say Kubric currently needs substantial compute for that reason §5. Converting ShapeNetCore.v2 took about 16 hours on an 80-core cloud VM §3.3.

Method

Framework

The framework has three parts: a Scene, a Camera and Assets §3.2. The Scene holds global settings such as resolution, frame count and gravity. Assets cover objects, lights, materials and animations. Adding an asset creates the matching object in each view, and the views stay in sync with the Kubric object §3.2.

• Simulation. The PyBullet backend runs rigid-body physics. It can also place objects without overlap. Its trajectories are turned into keyframes and a list of collision events §3, Simulator.
• Rendering. The Blender backend, driven through bpy, renders with Cycles, which models global illumination, reflection, refraction and subsurface scattering §3, Renderer.
• Annotations. From Blender's render passes the framework exports the RGB image plus depth, instance segmentation, optical flow, surface normals and object coordinates. It also records object metadata: 2D and 3D trajectories, bounding boxes, velocities, mass, friction, camera parameters and collision events §3, Annotations.
• Portability. Blender's Python module is hard to install and varies by version, so the authors ship a Docker image §3.1.

Worker pattern

A typical Kubric workflow writes one worker script that creates, simulates and renders one random scene. The dataset is produced by running the worker many times and collecting the output. Because scenes are independent and identically distributed, this scales from a local run to many cloud machines; the authors provide GCP launcher scripts §3.2. An Apache Beam pipeline that merges generation and collection is planned, not described as implemented §3.2.

The paper's example worker sets up a ground plane, one light and a camera, then adds eight KuBasic shapes. Each shape is placed by rejection sampling so that none overlap, given a random velocity, simulated in PyBullet and rendered in Blender §3.2, example-worker figure.

# Make one random scene and export its annotations  [§3.2]
# Input:  seed (random state), asset source, render settings
# Output: frames and object metadata written to disk

def make_scene(seed, asset_source, settings):
    # Create the random generator and the scene  [§3.2]
    rng = random_state(seed)
    scene = new_scene(settings.resolution, settings.frames, settings.fps)
    # The simulator and renderer both read the same scene  [§3.2]
    simulator = PhysicsSimulator(scene)
    renderer = PathTracer(scene, samples_per_pixel=settings.samples)
    # Static floor and light, plus a camera  [§3.2]
    scene += floor_object()
    scene += point_light(intensity=settings.light)
    scene.camera = random_camera(rng)
    # Add random objects, placed without overlap  [§3.2]
    for index in range(settings.count(rng)):
        obj = asset_source.create(choose(rng, asset_ids), random_material)
        obj.velocity = random_velocity(rng)
        scene += obj
        move_until_no_overlap(obj, simulator, rng=rng)
    # Run the physics, then render all passes  [§3.2]
    simulator.run()
    frames = renderer.render()
    # Export images and object metadata  [§3.2]
    write_image_dict(frames, output_dir)
    write_metadata(scene, output_dir)

Assets

• KuBasic. Eleven simple shapes for prototyping §3.3.
• ShapeNetCore.v2. A subset with about 51,300 models, canonically aligned and categorised. The authors removed doubles, disabled auto-smoothing, split sharp edges and slightly displaced faces, which reduced rendering artifacts. Collision shapes come from watertight meshes (ManifoldPlus), convex decomposition (V-HACD via PyBullet), and mass and inertia from trimesh. A few models failed conversion and are listed in the conversion code §3.3.
• Google Scanned Objects. About 1,000 textured household meshes, CC-BY 4.0, republished in Kubric format with collision meshes §3.3.
• Polyhaven. CC0 HDRI backgrounds and textures §3.3.

Assets are loaded from a manifest.json that lists archives, each with a JSON description of render and collision meshes and URDF physics §3.3.

SunDs

SunDs is a front-end on top of TensorFlow Datasets. Each dataset has a scenes split (scene-level metadata) and a frames split (per-example data). Tasks choose which features are decoded and which transform is applied, for example generating NeRF rays from camera parameters §3.4.

MOVi datasets (object discovery)

MOVi-A to MOVi-E are five video datasets of rising complexity. MOVi-A to -C have up to 10 moving objects over 24 frames. MOVi-D and -E have up to 23 objects, with only a few moving. MOVi-A uses CLEVR-style objects, MOVi-B adds more primitive shapes and colours, MOVi-C adds real-world backgrounds and scanned objects, and MOVi-E also moves the camera in a random direction §4.1. The MOVi-D and -E workers place static distractors first, let them settle, then toss dynamic objects movi_def_worker.py. The MOVi-A and -B workers spawn all objects in one region with a random velocity movi_ab_worker.py.

Point tracking baseline

The baseline adapts a self-supervised correspondence method from segment tracking to points §4.11.

# Track one query point through a video  [§4.11]
# Input:  frames (T, H, W, 3); query point (x, y) in frame t_query
# Output: estimated position and occlusion flag for each frame

def track_point(frames, query_xy, t_query, encoder):
    # Stride-8 features, L2-normalised per frame  [§4.11]
    features = [normalise(encoder(frame)) for frame in frames]
    # Feature of the query point by bilinear sampling  [§4.11]
    query_feature = bilinear(features[t_query], query_xy / 8)
    for t in range(num_frames):
        # Softmax over space of the similarity to the query  [§4.11]
        heat = softmax_2d(dot(query_feature, features[t]))
        # Keep only cells within 5 grid units of the peak  [§4.11]
        heat = zero_outside_radius(heat, argmax(heat), radius=5)
        # Estimate the location as the weighted mean of the heat  [§4.11]
        estimate[t] = weighted_mean_location(heat)
        # Map the estimate back to the query frame  [§4.11]
        back = cycle_back(estimate[t], t, t_query, features)
        # Occluded if the cycle ends over 48 px from the start  [§4.11]
        occluded[t] = distance(back, query_xy) > 48
    return estimate, occluded

The training loss is a contrastive term over spatio-temporal positions. Its weights come from bilinear interpolation of the ground-truth position and are zero where the target is occluded. The temperature is 0.1 §4.11. The paper writes it as:

\sum_{f_j \in F_i} \gamma_j \log\left(\frac{\exp(f^*_i \cdot f_j)/\tau}{\sum_k \exp(f^*_i \cdot f_k/\tau)}\right) §4.11

Other use cases

Each of the following is a dataset or evaluation design built with Kubric, with the setup the paper states:

• Optical flow (§4.2). Kubric scenes with 3D rigid motion are used to pretrain RAFT, compared with FlyingChairs and AutoFlow. The paper states the metrics, not the training schedule.
• NeRF texture and structure (§4.3). Flat surfaces with blue-noise textures at controlled frequencies, with a per-pixel frequency annotation.
• Pose estimation (§4.4). A MoveNet-style CenterNet trained with 80% real and 20% synthetic batches. Synthetic scenes use 41 rigged human models and motions from dance and workout animations.
• Pre-training (§4.5). A ResNet-50 is trained to classify ShapeNet objects pasted on random backgrounds, then transferred with a BiT-style protocol.
• Robust NeRF (§4.6). A "teleport" and a "jitter" variant add a moving impostor (a monkey head) to a static scene.
• Multi-view object matting (§4.7). Multi-view salient object detection with ShapeNet objects on HDRI backgrounds; an "easy" and a "hard" variant (clutter added).
• Complex BRDFs (§4.8). Highly specular ShapeNet renders with three lights and 12 ray bounces.
• Single-view reconstruction (§4.9). SoftRas trained on all of ShapeNetCore.v2 from 24 random viewpoints.
• Video-based reconstruction (§4.10). LASR fed with Kubric's ground-truth or estimated optical flow.
• Scene semantic segmentation (§4.12). KLEVR and ToyBox datasets compare 2D and 3D segmentation models.
• Novel view synthesis (§4.13). About one million scenes of ShapeNet objects on HDRI backgrounds, rendered at 128 x 128 with five conditioning and five target views.

Hyperparameters

Values stated for the datasets, with the source. "not stated" means the paper and the cited code do not give the value.

Results

Author-reported. Setups are as stated per section. Lower is better for AEPE and Chamfer distance. The novel-view-synthesis caption does not state the direction for LPIPS.

Object discovery (FG-ARI %, higher is better) §4.1

Optical flow, RAFT pretrained on each dataset §4.2

NeRF texture (NeRF trained on textured flat surfaces; setup not stated) §4.3

Pose estimation, keypoint mAP, COCO + Active training §4.4

Robust NeRF, test PSNR without impostors §4.6

Multi-view matting (salient object detection) §4.7. Metrics: max F-measure with β²=0.3, MAE, S-measure.

Complex BRDFs, few-view novel view synthesis §4.8

Single-view reconstruction (SoftRas) §4.9. Trained on all 55 ShapeNetCore.v2 categories. Best IoU: pillows 0.75, bowls 0.72. Worst: microphones and earphones 0.34 (in-distribution split, 80/20 per category). The per-category bar chart is in the paper's figure, not in its text.

Video-based reconstruction (LASR), Chamfer distance with true or estimated flow §4.10

Point tracking on MOVi-E, 256 x 256 §4.11

On vertically flipped videos the contrastive baseline reaches AJ 45.2, against 49.5 on the unflipped set §4.11.

Scene semantic segmentation (mean IoU) §4.12

Novel view synthesis on the 1M-scene dataset, 128 x 128 §4.13

Pre-training §4.5: the paper reports only a figure. Its text says the pilot "halves the gap" between random initialisation and ImageNet pre-training. No number is given in the text.

Paper vs code

Compared the official repository at commit 61f2422c84bab75006df33c6989e0b483db3ccfe (Apache-2.0) against the paper. I read the MOVi A/B/D/E worker scripts, the MOVi TFDS builder docstrings and the point-tracking dataset code. I did not check the other challenge workers.

• Point-tracking sampling cap. The paper caps samples per object at 0.2% of visible pixels §4.11. The code's max_sampled_frac default is 0.1 and applies to points on the sampling grid, so 0.1 means 10% dataset.py and dataset.py. The paper's experiments may have used another setting; this is not stated.
• Worker API. The paper's example worker builds a 640 x 640 scene and calls kb.Scene, kb.renderer.Blender and kb.write_image_dict directly §3.2, example-worker figure. The MOVi workers take CLI flags through kb.setup(FLAGS), use 256 x 256 by default and 64 samples per pixel movi_ab_worker.py. The figure is an illustration, not the MOVi configuration.
• Dataset count. The introduction and §4 describe 13 challenges §1; §4. The conclusions describe 11 case studies §5. The source also has commented-out notes that some challenges were moved to a supplement §4.
• Scene-count wording. The paper says the MOVi scenes have "up to 10" moving objects for MOVi-A to -C §4.1. The code matches (3 to 10) movi_ab_worker.py.
• ToyBox text vs table. The text says DeepLab outperforms NeSF "by 0.3% to 6.6%" in 2D §4.12. The table shows NeSF at 81.9 against DeepLab at 81.6 on ToyBox5, so NeSF is higher there §4.12 table.
• LASR text vs table. The text says ground-truth flow "improves performance, although this improvement is marginal" §4.10. For the airplane the gap is 0.62 against 1.52, and for walking the estimated-flow model is better (3.19 against 3.81).
• Physics values. The paper does not state friction, restitution or density for MOVi. The MOVi-A/B worker sets friction and restitution per material (metal 0.4 and 0.3, rubber 0.8 and 0.7) movi_ab_worker.py. The TFDS builder docstring gives densities of 2.7 for metal and 1.1 for rubber movi_a.py. The MOVi-D and -E workers keep the defaults of 0.5 and 0.5 for friction and restitution, per their docstring movi_d.py.

No discrepancies found in the licence: the repository's LICENSE is Apache-2.0, which matches §3.1.

Pitfalls

• Synthetic-to-real gap. The gains from synthetic training are measured on real benchmarks (Sintel, KITTI, COCO, Yoga) for optical flow and pose §4.2; §4.4. The other use cases are benchmarks or hypothesis tests, so do not assume transfer from them.
• Unstated training protocols. The optical-flow training schedule, the pre-training optimiser and the NeRF setup are not stated in the paper text §4.2; §4.3; §4.5. Rerunning the comparison needs the code or a re-derivation.
• Manual flow settings. The Kubric flow row uses "manually selected" rendering parameters, and the paper does not list them §4.2. The paper notes AutoFlow's parameters were learned for Sintel, so the comparison is not like-for-like §4.2.
• Physics and rendering are not free. Path tracing is much slower than rasterization §3.2. Blender's Python module varies by version, so use the Docker image §3.1.
• Depth and flow encoding in MOVi. Depth and flow are stored as uint16 and rescaled by the range recorded in the metadata movi_a.py. Converting without that range gives wrong values.
• Object metadata. MOVi segmentation IDs are one higher than the index in the instance list. The TFDS docstring warns about this movi_a.py.
• Dataset size. The TFDS builders describe "approx 10k videos" per MOVi dataset movi_a.py. The paper does not state the number of MOVi videos.
• Asset licences. The paper says non-CC0 asset licences prevent publishing the pose-estimation data, which can be regenerated only by someone who owns the same assets §4.4. The paper does not state ShapeNet's licence terms, so check them before redistributing generated data.
• Point-tracking occlusion. The occlusion check is worse than chance at the 48 px threshold, and the authors drop failed points to raise AJ §4.11. Scores are sensitive to that choice.
• Robust NeRF wording. The text says unstructured outliers are addressed "to some degree" by NeRF-L1 at "−2.4 dB" §4.6. The table's only pair that differs by 2.4 dB is NeRF-L1 static against teleport (42.8 against 40.4), so the text probably means that, but it does not say which pair.
• Fine-tuned SOD columns. The text says the best models were fine-tuned on the easy set and then evaluated on the hard set §4.7. The table does not say which training set each column used.
• Single-view and video reconstruction. Both are evaluations of existing methods on new data, not new methods. The LASR results show it fails on a torus, which has non-zero genus §4.10.
• Scale claims. The paper's "TB" data sizes §1; §4 and the cluster runs are not checked against the repository; the code shows only the worker pattern.

Lineage

• Builds on: PyBullet §3.2 and Blender's Cycles renderer §3.2. Object discovery compares SAVi and SIMONe §4.1. Optical-flow comparisons use FlyingChairs and AutoFlow §4.2. The 3D reconstruction work uses SoftRas and LASR §4.9; §4.10. Novel view synthesis compares LFN, pixelNeRF and SRT §4.13; see pixelNeRF.
• Compared against: general pipelines ThreeDWorld, UnrealCV, iGibson, Habitat, BlenderProc, and real-to-sim pipelines OpenRooms and Omnidata §2.
• Superseded by: not known from content in this wiki.