Foundation Physics Graph Model — object velocity extraction

Extracts metric 3D world-frame velocity (m/s) of the object a robot is manipulating, from DROID episodes:

DROID episode mp4 ──► SAM 3.1   (text prompt derived from the episode's task instruction)
                  ──► mask of the manipulated object
                  ──► query points sampled inside the mask
                  ──► TAPNext++ online point tracking
                  ──► lift 2D tracks to metric 3D  (depth from PointWorld scene_flows)
                  ──► camera extrinsics ⇒ robot-base world frame ⇒ velocity (m/s)

The output is intended as node/edge features of a physical scene graph, so every velocity is metric, timestamped, and carries explicit reference-frame provenance (VelocityEstimate.frame).

---

Why two data sources

nvidia/PointWorld-DROID contains no RGB video — only a single initial_rgb JPEG per clip. It is a derived-annotation release (~3.91 TB packaged). So the pipeline joins two sources on the DROID episode uuid (AUTOLab+0d4edc83+2023-10-21-19h-07m-04s):

droid/depth_320x180/ is a 1.23 TB archive that cannot be cherry-picked, so it is never downloaded. scene_flows — which is already (T, N, 3) metric 3D point tracks — serves as the depth source instead.

Total footprint for a 3–5 episode run is roughly 12 GB including both model checkpoints.

---

Setup

1. GPU driver shim (required)

This host's NVIDIA kernel module is 570.172.08 while the userspace libraries were upgraded to 580.173.02, so a bare nvidia-smi fails with Failed to initialize NVML: Driver/library version mismatch. The 570 libraries are still on disk, so a symlink shim fixes it without sudo or a reboot:

source scripts/nvidia_lib_shim.sh
nvidia-smi          # should now list the H200s

Everything that touches the GPU must run with that shim active. fpgm.utils.gpu.ensure_cuda() fails fast and repeats this remedy if it is missing.

2. Environment

bash scripts/setup_env.sh
conda activate fpgm

This creates a Python 3.12 conda env (SAM 3.1 requires ≥3.12) and installs torch 2.10.0 from the cu128 index. That index is not optional: driver 570 caps CUDA at 12.8, and the default PyPI wheels are now cu13x, which need driver ≥580 and will fail at runtime here.

It also clones sam3 and tapnet into third_party/ (pinned commits recorded in third_party/PINNED_COMMITS.txt) and installs them editable. tapnet is installed with --no-deps on purpose: its pyproject.toml lists jax, dm-haiku, jaxline and optax as unconditional dependencies, none of which the TAPNext++ PyTorch path ever imports.

3. Data and checkpoints

export HF_TOKEN=<your token>            # facebook/sam3.1 is a gated repo

python scripts/download_checkpoints.py   # SAM 3.1 (3.5 GB) + TAPNext++ (2.5 GB)
python scripts/download_pointworld.py --shard shard-000000 --yes
python scripts/download_droid_episodes.py --from-shard shard-000000 --limit 5

Checkpoints are resolved through the standard HuggingFace cache, which on this host already holds 71 GB — nothing is duplicated into the repo.

4. Run

source scripts/nvidia_lib_shim.sh
python scripts/run_pipeline.py --config configs/droid_velocity.yaml

Per-clip overlay videos, contact sheets and run_summary.json land in outputs/.

---

Layout

src/fpgm/
  types.py            shared dataclasses — the contract every stage speaks
  config.py           dataclass config schema + YAML loader
  data/               episode identity, DROID raw client, PointWorld store
  prompting/          task instruction -> SAM 3.1 noun phrase
  segmentation/       VideoSegmenter ABC + SAM 3.1 adapter
  tracking/           PointTracker ABC + query sampling + TAPNext++ adapter
  depth/              DepthSource ABC + scene-flow depth
  geometry/           camera, SE3, convention detection, lifting, velocity, validation
  pipeline/           clip frame extraction + the orchestrator
  viz/                mask / track / velocity overlays

Stages communicate only through the dataclasses in types.py; no stage reaches into another's internals. That is what makes each one testable against synthetic data with neither a GPU nor a download:

pytest -q

---

Design notes worth knowing before changing things

Camera resolution is a typed property, not a convention. CameraIntrinsics carries the (width, height) it is valid at, and Camera.project/unproject never take a resolution argument. The h5 intrinsic belongs to the annotation resolution while mp4 frames are a different size, so Camera.rescaled() is the single, grep-able place where resolution changes.

The scene-flow frame convention is detected, never assumed. Whether scene_flows are stored world- or camera-frame is undocumented. Each clip is tested empirically by reprojecting into initial_rgb and scoring in-bounds fraction plus colour agreement. An indecisive margin raises AmbiguousConventionError; both hypotheses scoring badly raises CameraCalibrationMismatchError (a different fault — wrong intrinsic/extrinsic/ serial pairing). Guessing here would apply a systematic rigid transform to every 3D position while still looking entirely plausible.

Depth neighbours are restricted to the object mask. Depth at a tracked point is interpolated from projected scene_flows points, but only from those inside the same SAM 3.1 mask. A background point that is near in image space sits on a different surface, and letting it into the interpolation is the dominant failure mode when tracking a moving object. Delaunay/barycentric interpolation was considered and rejected: its triangles silently span occlusion boundaries, producing smoothly wrong depth with no signal that anything went wrong.

Long occlusions are not bridged. Gaps up to max_gap_frames are interpolated in position before differentiating. Longer gaps split a track into independent runs with NaN across the gap, because a long occlusion is exactly where the object may have changed direction.

Nothing is invalid without a reason code. DepthResult.method records why a point was rejected. Debugging "why is this velocity NaN" three stages downstream without provenance is the single biggest time sink in a pipeline like this.

The GPUs are shared. SAM 3.1 holds every frame as a (T, 3, 1008, 1008) fp16 tensor (~6.1 MB/frame), so offload_video_to_cpu defaults to true and sessions are always closed in a finally. use_fa3 defaults to false because FlashAttention-3 is a separate install.

---

Validation

scene_flows already contains ground-truth 3D motion for annotated points inside the mask, so accuracy is checked without any external source: a Track3D built directly from scene_flows is run through the same VelocityEstimator and compared against the pipeline's own output. Magnitude and direction (cosine) errors are reported separately, because a magnitude error implies a depth-scale bug while a direction error implies a convention/axis bug — the report tells you which stage to fix, not merely that something is wrong.

---

Licences

DROID and PointWorld-DROID carry their own terms (PointWorld is NVIDIA research-only); SAM 3.1 is under the SAM License and its HuggingFace repo is gated with manual approval. This repository does not redistribute any of them.