twanghcmut's picture
|
download
raw
7.85 kB

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:

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.

Xet Storage Details

Size:
7.85 kB
·
Xet hash:
cefc839430a8c71423052ad9b75d191c3cb92092dafeb248dea27037f11fe040

Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.