README.md

# 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`):

| Need | Source | Cost |
| --- | --- | --- |
| RGB video + **task instruction** | DROID raw on GCS, anonymous HTTPS | ~16 MB/episode |
| 3D point tracks, intrinsics, extrinsics | PointWorld `*_flows.h5` | 3.5 GB (1 shard ≈ 62 episodes) |
| Camera extrinsics for **every** episode | PointWorld `droid/cameras/` | 19.5 MB (all 42,935) |

`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**:

```bash
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
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

```bash
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

```bash
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:

```bash
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.