Skip to content

feat(arc): arc tokenize and detokenize as graph nodes - #11

Open
AnikethCheluva wants to merge 15 commits into
graph-hpt-configsfrom
graph-arc
Open

feat(arc): arc tokenize and detokenize as graph nodes#11
AnikethCheluva wants to merge 15 commits into
graph-hpt-configsfrom
graph-arc

Conversation

@AnikethCheluva

Copy link
Copy Markdown

Port of GaTech-RL2#573's arc-tokenizer work onto the graph fork, reduced to the part that
belongs in this architecture: the tokenizer as a pair of pipeline nodes.

The Planar arc tokenizer already existed as a loader-side transform, and that
placement is invisible to the graph -- a stage list shows the model consuming
target with no sign the target is an arc token rather than a time-indexed
chunk, and tools/config_graph.py cannot lint the boundary because nothing
declares it. These two stages move the boundary into the graph:

ArcTokenizeStage takes ActionTargetBuilder's PLACE rather than following it.
Two writers of target would be a duplicate-writer lint error and would leave
the graph ambiguous about which one the denoiser models; as the sole writer, it
leaves the rest of a DP or flow chain untouched. It tokenizes per sample through
the same TokenizePlanarArcLength the loader transform uses, so both paths
produce identical targets -- there is a test pinning that.

ArcDetokenizeStage walks the waypoint polyline at speed * dt * k, the
inverse of how the tokenizer lays waypoints out uniformly in arc length. It
saturates at the window end rather than extrapolating past it, holds position
for a zero-speed token (matching the tokenizer's degenerate branch), and
interpolates heading as (cos, sin) so a chunk that wraps past +/-pi does not
unwind through zero.

Adds inference_only to core as the mirror of train_only. It is needed, not
cosmetic: this runner treats a stage with an unsatisfied read as a
configuration error and raises, so a stage that legitimately exists in only one
mode has to say so. Without it the detokenizer -- whose pred_action input
exists only at inference -- would abort every training run. config_graph.py
reports the new restriction alongside the existing one.

Co-Authored-By: Claude Opus 5 noreply@anthropic.com
Claude-Session: https://claude.ai/code/session_01DcPCNAW3GKERANrqvfSCY1

@AnikethCheluva
AnikethCheluva marked this pull request as ready for review September 5, 2026 21:01
@ElmoPA
ElmoPA changed the base branch from graph-abc to graphite-base/11 September 5, 2026 23:05
@ElmoPA
ElmoPA changed the base branch from graphite-base/11 to main September 5, 2026 23:05
@AnikethCheluva
AnikethCheluva changed the base branch from main to graphite-base/11 September 6, 2026 20:05
@AnikethCheluva
AnikethCheluva changed the base branch from graphite-base/11 to graph-abc September 6, 2026 21:29
@ElmoPA
ElmoPA changed the base branch from graph-abc to graphite-base/11 September 7, 2026 09:03
@ElmoPA
ElmoPA changed the base branch from graphite-base/11 to main September 7, 2026 09:03
@AnikethCheluva
AnikethCheluva changed the base branch from main to graphite-base/11 September 7, 2026 14:28
@AnikethCheluva
AnikethCheluva changed the base branch from graphite-base/11 to graph-abc September 7, 2026 14:28
@ElmoPA
ElmoPA changed the base branch from graph-abc to graphite-base/11 September 7, 2026 17:21
@ElmoPA
ElmoPA changed the base branch from graphite-base/11 to main September 7, 2026 17:22
@AnikethCheluva
AnikethCheluva changed the base branch from main to graphite-base/11 September 8, 2026 05:01
@AnikethCheluva
AnikethCheluva changed the base branch from graphite-base/11 to graph-hpt-configs September 8, 2026 05:02
@AnikethCheluva
AnikethCheluva changed the base branch from graph-hpt-configs to graphite-base/11 September 8, 2026 05:02
AnikethCheluva and others added 15 commits September 10, 2026 18:18
Port of GaTech-RL2#573's arc-tokenizer work onto the graph fork, reduced to the part that
belongs in this architecture: the tokenizer as a pair of pipeline nodes.

The Planar arc tokenizer already existed as a loader-side transform, and that
placement is invisible to the graph -- a stage list shows the model consuming
`target` with no sign the target is an arc token rather than a time-indexed
chunk, and tools/config_graph.py cannot lint the boundary because nothing
declares it. These two stages move the boundary into the graph:

`ArcTokenizeStage` takes ActionTargetBuilder's PLACE rather than following it.
Two writers of `target` would be a duplicate-writer lint error and would leave
the graph ambiguous about which one the denoiser models; as the sole writer, it
leaves the rest of a DP or flow chain untouched. It tokenizes per sample through
the same TokenizePlanarArcLength the loader transform uses, so both paths
produce identical targets -- there is a test pinning that.

`ArcDetokenizeStage` walks the waypoint polyline at `speed * dt * k`, the
inverse of how the tokenizer lays waypoints out uniformly in arc length. It
saturates at the window end rather than extrapolating past it, holds position
for a zero-speed token (matching the tokenizer's degenerate branch), and
interpolates heading as (cos, sin) so a chunk that wraps past +/-pi does not
unwind through zero.

Adds `inference_only` to core as the mirror of `train_only`. It is needed, not
cosmetic: this runner treats a stage with an unsatisfied read as a
configuration error and raises, so a stage that legitimately exists in only one
mode has to say so. Without it the detokenizer -- whose `pred_action` input
exists only at inference -- would abort every training run. config_graph.py
reports the new restriction alongside the existing one.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01DcPCNAW3GKERANrqvfSCY1
…path

Ports `fix(arc): detokenize walked arc length at a chord rate` (df5b849) from
the main repo's arc branch. Same defect class, different pair of metrics.

There, the velocity token was a CHORD rate (net displacement over time) walked
against cumulative ARC positions. Here, TokenizePlanarArcLength derives its
speed from `end`, a cumulative SE(2) length that already includes
`lambda * rotation` -- so the token's speed is an SE(2) rate. ArcDetokenizeStage
accumulated translation alone and advanced through it at that rate, which
traverses the window too fast by exactly the ratio between the two lengths.

That ratio is unbounded for a rotating path. A 60-unit translation carrying
2.5 rad at radius 30 measures 135 in SE(2): the reconstruction ran 2.25x fast
and saturated after 18 of the 39 steps the motion actually took, piling the rest
of the horizon on the endpoint. The shipped experiment config had
`rotation_radius: ${planar.arc_rotation_radius}` = 30, so this was live, and
nothing in the tensor shapes showed it. Translation-only was correct only at
radius 0 -- which is exactly the case the original tests covered.

The decoder now accumulates the same metric as `planar_step_distance`, taking
the angular step as the principal difference via atan2 rather than unwrapping a
decoded angle, so the +/-pi seam cannot trip it. `rotation_radius` is a
constructor argument and the config points both nodes at one `planar.*` value; a
test asserts the two agree on it, along with dt and waypoint count.

The other two arc commits need no code here. `feat(arc): tokenize the raw window
directly` (e182253) fixes a loader-side InterpolatePose that decimated the
window before tokenizing and left dt too large; tokenizing in the graph has no
such step -- the dense transform only pads -- so this fork already had the fixed
behaviour, and there are now tests pinning it rather than leaving it incidental.
`fix(arc): detokenize before the camera-frame revert` (56f9d67) is about a
frame revert applied to all rows of a token, translating the velocity row into a
position; the Planar arc decoder only ever reads waypoint zero and the
detokenize node slices the timing row out explicitly, so the hazard has no path
into this fork.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01DcPCNAW3GKERANrqvfSCY1
…aining recipes

Bundles three related additions built on top of the arc tokenizer stack.

BimanualCartesianEval + eval_bimanual_cartesian.yaml + viz/cartesian/base.yaml
add a validation-time evaluator that decodes both arms' cartesian actions and
overlays predicted vs GT waypoints on the front camera. Used by the abc_arc
experiments below as the default `evaluator`.

New PACE launcher configs cover the non-blackwell partitions:
submitit_pace_l40s, submitit_pace_h100, submitit_pace_a100,
submitit_pace_rtx6000, submitit_pace_v100. Each hardcodes `gres:
"gpu:<type>:2"` for the 2-GPU launch pattern (this repo does not register an
`eval` OmegaConf resolver, so the `${eval:...}` arithmetic used by
submitit_pace/submitit_skynet does not compose), and respects the per-partition
CPU:GPU cap (4:1 on l40s/h100, 6:1 on rtx6000, 8:1 on a100, 12:1 on v100).
mem_per_gpu is dialed per card size. The same `${eval:...}` gres bug is fixed
in submitit_pace.yaml (H200) by hardcoding `gpu:h200:2`.

abc_arc data/model/experiment configs add the training recipes: the fstshirt
BC baseline (arc + time-indexed twin), the fstshirt mecka freefold cotrain
(arc D40 M100 + baseline), and the mecka fold multitask cotrain (arc D40 M100
+ baseline). Each experiment picks the matching data + model pair and points
at eval_bimanual_cartesian. 300M model variants ship the wider
`down_dims: [512, 1024, 2048]` DP backbone for the cotrain.

Also gitignores .claude/ so local Claude Code settings do not get committed.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
- PipelineAlgo._move_value: cast float64 tensors to float32 on move so
  bf16 AMP autocast doesn't hit dtype mismatch on Linear (zarr stamps
  ee_pose as float64).
- BimanualCartesianEval: port to main-repo EvalVideo format — one h264
  mp4 per episode (buffered on episode_hash boundary, chunked fallback),
  rank-0 only, plus wandb.Video upload of every finished mp4 in
  on_validation_end. Drops per-batch PNG + wandb.Image path.
- eval_bimanual_cartesian.yaml: replace videos_per_epoch with
  video_chunk_frames/max_episode_frames; default video_output_dir=null
  so writes land under {default_root_dir}/videos.
Two things that stopped every abc experiment config from running.

override hydra/launcher needs a leading slash. Without it hydra resolves the
group relative to the experiment package and aborts with "Could not override
'experiment/hydra/launcher@hydra.launcher'. No match in the defaults list", so
the configs failed to compose at all, on PACE as well as skynet.
abc_mecka_fold_multitask_cotrain_baseline already had the right form.

The two baseline configs asked for action_horizon 16 while the loader hands
over 100-row chunks: Yam.get_transform_list interpolates to chunk_length,
which defaults to 100, and the data configs never override it. The diffusion
stages validate target width, so training died on the first batch with
"Diffusion target must be (32, 16, 14), got (32, 100, 14)". 100 matches the
sibling yam_mecka_cotrain_dp_baseline model, which hardcodes it in all three
diffusion slots. The arc configs are untouched: their 16 is the detokenizer
output horizon and their diffusion stages read arc_token_rows instead.

Verified by a debug run on skynet (2080_ti, overcap): reaches global_step 20
with populated optimizer state. Note trainHydra catches exceptions and exits
0, so slurm reported COMPLETED for the two runs that trained nothing.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Ports arc_length_tokenizer.py from the main repo's arc branch. The arc data
configs synced in c35d057/bd3659b reference this module, so abc_fstshirt_arc_bc
could not instantiate its dataset without it.

Turns a (T, 14) bimanual cartesian chunk into (M+1, 14): M waypoints uniform in
each arm's arc length over the first D meters of travel, plus a velocity token
in row M. xyz and gripper are linear-interpolated at the arc-length targets, ypr
is slerped so orientation stays supervised.

Self-contained, only numpy and scipy. Includes the chord-vs-arc detokenize fix:
the velocity token is a chord rate but the traversal walks arc length, and arc
>= chord off a straight line, so replay ran up to 1.98x slow on real folding
data before the conversion.

This is the loader-side bimanual tokenizer, distinct from the Planar graph-node
tokenizer in stages_arc.py -- that one handles 5-wide SE(2) tokens for pushshapes
and stays as is.

Brings the module's 53 tests, passing unmodified.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
1ced546 changed PipelineAlgo._move_value to cast float64 to float32 on device
move, but left this test asserting the old dtype-preserving behaviour, so the
branch was red. Updates the assertion and the name to describe what the move
now does, keeping the integer-metadata half that still holds.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Completes the arc path for abc_fstshirt_arc_bc, which could not instantiate its
dataset: the data config asks Yam for action_mode=arc_tokenizer_cartesian with
D/M args and the fork only had plain cartesian.

Yam gains the arc action mode plus the keymap's wider raw window
(ARC_TOK_ACTION_HORIZON=200 vs 45). Arc defaults chunk_length to that raw window
so nothing resamples before tokenizing -- interpolating to 100 first would
decimate the window 2x and arc length measured on a decimated path reads short.
_append_arc_tokenizer lands in eva.py, matching upstream's layout so future
syncs don't conflict.

The evaluator is the part that needs care. An arc token's rows 0..M-1 are
waypoints but row M is a VELOCITY, and the revert transforms rotate and
translate every row they get, so a raw token comes back with that velocity
turned into a position. The overlay still renders, which is what makes it
expensive: the trajectory looks plausible while the arm appears to teleport.

So BimanualCartesianEval gets a _viz_source seam (identity -- a time-indexed
run already predicts poses) and ArcBimanualCartesianEval overrides it to
detokenize, meaning apply_transform only ever sees rows that really are poses.
Metrics stay in token space, where prediction and target are both tokens.

On real ABC data the velocity row's magnitude is 0.169 against a waypoint's
0.157, so nothing about its scale flags it -- a test pins that, since it is the
reason the ordering matters. Others pin that eval D/M match the data
tokenizer's (a mismatch replays at the wrong speed with every shape still
lining up) and that the baseline config does not pick up the arc evaluator.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
The token carried ONE mean arc speed for the whole chunk, which only describes
constant-speed motion. Adds velocity_mode as an opt-in alternative that keeps a
local rate per waypoint interval, so a chunk that accelerates, decelerates or
dwells replays at its original pace.

  mean          M waypoints + one mean-speed row   -> M + 1 rows (default)
  per_waypoint  M waypoints + M local-rate rows    -> 2 * M rows

Each rate comes from the elapsed time at the sampled arc positions, recovered
from the bracketing source frames as (index + alpha) * dt rather than by
dividing the window by one duration. That is what preserves local speed changes
and stationary frames before motion, both of which a chunk mean erases.
Decoding integrates the rates into a cumulative time curve and inverts it, so
control steps land where they did in the source.

Round-trip error against the recorded actions, M=32 D=200 R=30deg:

  trajectory         mean xy   granular xy   mean ang   granular ang
  constant              0.00          0.00     0.0000         0.0000
  accelerating          4.80          0.28     0.1277         0.0000
  decelerating         23.38          0.15     0.0648         0.0000
  dwell_then_move       0.72          0.00     0.2912         0.0001
  stop_and_go           1.67          0.00     0.0110         0.0000

Identical where a mean rate is a complete description, up to 150x better where
it is not. Recovering the time parameterization fixes heading too -- rotation is
sampled at the same arc positions, so correct timing lands it in the right
place without decoupling the streams.

Modular by design: planar_arc.arc_token_rows is the single source of truth for
the layout, and the tokenizer, both graph stages, the native decoder and the
configs all size themselves from it. One planar.arc_velocity_mode knob feeds
the tokenize node, the detokenize node and the evaluator's decoder, with tests
asserting all three read it and that arc_token_rows tracks the mode.

mean stays the default: the token width differs between modes, so existing
checkpoints and norm stats do not transfer. Both modes lint clean in train and
inference.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Extends velocity_mode to the bimanual (ABC/yam) tokenizer, which hardcoded
MEAN_PER_DIM. Its PER_STEP_* enum values existed but detokenize rejected them,
so the granular option was unreachable on this path.

  mean          M waypoints + one velocity row  -> M + 1 rows
  per_waypoint  M waypoints + M velocity rows   -> 2 * M rows

Each interval's rate is the waypoint delta over the elapsed time across that
interval, with the time recovered from the SOURCE frames by bracketing the
waypoint's arc position in the arm's cumulative arc length. Detokenize
integrates those durations into a cumulative time curve, inverts it to an arc
position per control step, and reuses the existing interpolators.

Worth noting the chord-vs-arc correction the mean path needs does not apply
per-interval: consecutive waypoints are joined by straight segments, so an
interval's chord IS its arc length. That discrepancy only arises for one rate
spanning the whole token, which is what made the mean mode wrong.

Round-trip xyz error against the recorded chunk, metres, D=0.40 M=32:

  trajectory          mean      granular
  constant         0.00151       0.00003
  accelerating     0.01727       0.00027
  decelerating     0.08377       0.00019
  dwell_then_move  0.01130       0.00020

8.4cm -> 0.02cm on a decelerating chunk. Even constant improves, because the
mean path's chord-to-arc ratio is approximate where per-interval is exact.

One knob, abc.arc_velocity_mode, feeds the data transform, the evaluator and
the model's row count; bimanual_arc_token_rows is the single source of truth
for the layout and tests assert all sites agree. abc_fstshirt_arc_bc switches
to per_waypoint, which also makes the token 200 rows -- divisible by four, so
it clears the ConditionalUnet1D skip-connection crash that 101 odd rows hit.

Verified end to end on skynet: global_step 4, val mp4s written with overlays
drawn, Valid/MSE and Valid/Native_MSE finite.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Brings the arc metric families over from the main repo's arc branch, where they
lived inside an HPT-coupled eval class. The metrics are pure numpy over
(T, 14) chunks, so they land in egomimic/eval/arc_metrics.py with no evaluator,
trainer or model coupling -- testable without a graph.

Four families, all scored against the same time-indexed ground truth:

  arcmatch    re-tokenize both sides onto a matched per-arm span and score the
              waypoints. Travel is divided out, so this is path SHAPE alone.
              Reported with and without the velocity row; the gap is timing.
  dtw         warp the prediction against the GT chunk. Elastic in time, so
              unlike arcmatch it DOES see a travel mismatch.
  chunk       plain MSE by component: xyz, ypr, gripper, final.
  pose_err_m  position and rotation as one number in metres, via a lever arm,
              so radians become the length they cost rather than being summed
              with metres under an arbitrary unit choice.

Both arms of the ablation use the SAME evaluator and land on the same charts,
reaching the metrics by their own paths. The shared scoring lives on
BimanualCartesianEval behind two hooks: an arc run detokenizes its token into
control steps, a baseline run already predicts poses and de-interpolates to the
raw window instead. _is_arc detects which from the prediction's shape rather
than being configured twice.

De-interpolation is not cosmetic. The transform list interpolates the raw yam
window up to chunk_length (100), and arc length on an interpolated path
chord-cuts between inserted samples and reads systematically short, so arc
scoring runs on the raw 45 rows. Indices are selected rather than
re-interpolated, so no samples are invented and rotations are never
interpolated twice.

The ground truth comes from actions_cartesian_untokenized, which the tokenizer
preserves precisely because recovering it by detokenizing would be circular: a
reconstruction spans D by construction, so its distance says nothing about how
far the arm actually travelled.

Verified rather than assumed. The matched-span rule was checked against the
spec written out literally -- branch per arm on whether the GT or prediction
travelled less, cut whichever is longer -- across six cases (equal, 2x, 0.5x,
either side cut short, different shape), bit-identical every time: both
branches reduce to min(gt travel, pred travel) per arm. A perfect prediction
scores exactly 0.0 on every family. Travel invariance holds exactly -- same
path at 0.5x/0.75x/1.0x travelled gives arcmatch_xyz_mse 0.00e+00 while
dtw_xyz_l2_m rises 0.0000 -> 0.0198 -> 0.0786. pose_err_m matches its closed
form for pure rotation and pure translation, and the geodesic is invariant
under a shared frame change, which a per-axis ypr difference is not.

Metric tensors are built on the prediction's device: log_dict(sync_dist=True)
all-reduces them and NCCL has no CPU support, so a host-side numpy result cast
to a CPU tensor raises "No backend type associated with device type cpu" the
moment a GPU run syncs.

D, M, arcmatch_points and arc_chunk_rows are pinned equal across the two
experiments, with a test asserting it, because a drift there silently makes the
arms incomparable while every number still looks plausible.

Read arcmatch_span_m and arcmatch_travel_ratio beside the arcmatch scores: with
travel divided out, a policy that stalls scores well on shape while those fall.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Ports Human's arc path from the main repo's arc branch, the counterpart of the
yam work. Both cotrain arc configs failed outright before this: they ask Human
for arc_tokenizer_cartesian_gripper_padded and the fork only had plain
cartesian.

Three pieces, all mirroring what yam needed:

- ARC_TOK_ACTION_HORIZON = 600 for the arc keymap. Human is subsampled by
  stride, so at stride=3 those 600 raw frames yield 200 samples -- matching
  yam's arc window in physical time (~6.7 s) rather than in row count. The arc
  keymap is otherwise byte-identical to cartesian, so one transform list serves
  both, and a test asserts that.
- chunk_length defaults to the raw window for arc, i.e. no resampling before
  tokenizing. Interpolating to 100 first would decimate the human window and
  arc length on a decimated path reads systematically short.
- Gripper padding runs BEFORE the tokenizer. Human has no gripper signal and
  the tokenizer's layout routes gripper into slot 6 per arm, so the zero column
  has to exist by then. Bare arc_tokenizer_cartesian is rejected with a pointer
  to the padded mode rather than failing at the first batch on a 12D chunk.

The detail worth calling out is dt. The chunk is subsampled by
actions[::stride], so consecutive samples are stride/30 s apart, not 1/30.
Leaving the tokenizer's default inflates the velocity channel by exactly
stride. It cancels inside tokenize -> detokenize, so a round-trip test would
not catch it, but it is what the model learns and what a deployed policy would
command. Yam is unstrided and keeps 1/30; a test pins both.

Also unblocks the two cotrain arc configs for a second reason. Their models
hardcoded action_horizon 101 -- M + 1, odd -- which crashes
ConditionalUnet1D's skip-connection concat, since two downsample stages need
the row count divisible by four. They now read abc.arc_token_rows and default
to per_waypoint, giving 2 * M = 200, and a test runs a real tensor through the
denoiser at that width rather than trusting the arithmetic.

Verified on real data: a mecka episode and an abc episode both tokenize to
(200, 14) with human's gripper columns all zero and yam's populated, so a
cotrain batch is layout-consistent. Human's preserved window is (600, 14) and
yam's (200, 14), as their horizons imply.

Not ported: arc_tokenizer_keypoints, which needs keypoint_arc_tokenizer.py and
is a different action space no config here selects.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant