feat(arc): arc tokenize and detokenize as graph nodes - #11
Open
AnikethCheluva wants to merge 15 commits into
Open
feat(arc): arc tokenize and detokenize as graph nodes#11AnikethCheluva wants to merge 15 commits into
AnikethCheluva wants to merge 15 commits into
Conversation
This was referenced Sep 5, 2026
Author
|
Warning This pull request is not mergeable via GitHub because a downstack PR is open. Once all requirements are satisfied, merge this PR as a stack on Graphite.
This stack of pull requests is managed by Graphite. Learn more about stacking. |
AnikethCheluva
marked this pull request as ready for review
September 5, 2026 21:01
ElmoPA
force-pushed
the
graphite-base/11
branch
from
September 5, 2026 23:05
67bdd81 to
3d60f06
Compare
ElmoPA
force-pushed
the
graphite-base/11
branch
from
September 7, 2026 09:03
67bdd81 to
3d60f06
Compare
AnikethCheluva
force-pushed
the
graphite-base/11
branch
from
September 7, 2026 14:28
75525eb to
67bdd81
Compare
ElmoPA
force-pushed
the
graphite-base/11
branch
from
September 7, 2026 17:22
67bdd81 to
3d60f06
Compare
AnikethCheluva
force-pushed
the
graph-arc
branch
from
September 8, 2026 05:02
bbe3a73 to
3d749b6
Compare
AnikethCheluva
changed the base branch from
graphite-base/11
to
graph-hpt-configs
September 8, 2026 05:02
This was referenced Sep 8, 2026
AnikethCheluva
changed the base branch from
graph-hpt-configs
to
graphite-base/11
September 8, 2026 05:02
AnikethCheluva
force-pushed
the
graph-arc
branch
from
September 8, 2026 21:07
3d749b6 to
b7c4ccf
Compare
AnikethCheluva
force-pushed
the
graphite-base/11
branch
from
September 8, 2026 21:07
118f0b2 to
c142158
Compare
AnikethCheluva
changed the base branch from
graphite-base/11
to
graph-hpt-configs
September 8, 2026 21:07
AnikethCheluva
force-pushed
the
graph-hpt-configs
branch
from
September 9, 2026 22:21
c142158 to
61eee62
Compare
AnikethCheluva
force-pushed
the
graph-arc
branch
from
September 9, 2026 22:21
b7c4ccf to
5ddc150
Compare
AnikethCheluva
force-pushed
the
graph-arc
branch
from
September 10, 2026 17:56
5ddc150 to
e35594b
Compare
AnikethCheluva
force-pushed
the
graph-arc
branch
from
September 10, 2026 18:07
1ab87e9 to
70207bc
Compare
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>
AnikethCheluva
force-pushed
the
graph-arc
branch
from
September 10, 2026 22:20
70207bc to
80a84f2
Compare
AnikethCheluva
force-pushed
the
graph-hpt-configs
branch
from
September 10, 2026 22:20
2a2d247 to
111cb46
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.

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
targetwith no sign the target is an arc token rather than a time-indexedchunk, and tools/config_graph.py cannot lint the boundary because nothing
declares it. These two stages move the boundary into the graph:
ArcTokenizeStagetakes ActionTargetBuilder's PLACE rather than following it.Two writers of
targetwould be a duplicate-writer lint error and would leavethe 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.
ArcDetokenizeStagewalks the waypoint polyline atspeed * dt * k, theinverse 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_onlyto core as the mirror oftrain_only. It is needed, notcosmetic: 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_actioninputexists 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