Push T HPT training on freeform preference rankings - #532
Draft
RyanPCo wants to merge 37 commits into
Draft
Conversation
Introduces MultiDataset as the dataset graph wrapper that owns per-embodiment NormStats, runs bounds-check + per-sample retry in __getitem__, and applies normalization after the leaf returns. Stats are shared across nested MultiDatasets by reference via set_norm_stats_from. Bounds-check is implemented as a method on MultiDataset, not as a per-leaf transform — this avoids the shared-transform-list aliasing trap where appending Reject+Normalize once per leaf would accumulate N pairs on a single shared list (the resolver passes one transform_list reference to every leaf). The leaf's transform pipeline stays purely embodiment-side. Retry on bounds/NaN/Inf violation resamples globally within the failed sub-dataset (not per-episode), and runs as a bounded while-loop so a systematically bad slice raises a clean error instead of blowing Python's recursion limit.
…e obsolete data_schematic configs
The data layer no longer needs paligemma/control-mode/proprio knowledge. MultiDataModuleWrapper now only stacks tensors and preserves list-valued keys; PI.process_batch_for_training samples a prompt per item from the raw annotation lists, splices in the optional Embodiment/Control-mode/ State blocks, and tokenizes. All prompt knobs moved to pi0.5_base.yaml (override per-run via model.robomimic_model.*); 6 data configs lost their prompt blocks (would otherwise crash on the new __init__). build_tokenized_collate is kept in pl_data_utils for rollout.py and the deprecated wrappers that still call it.
- New tools: visualize_episode (replay obs+action), scripted_collect (headless heuristic), dataset_stats (per-config summary). - Physics: smoother kinematic pusher (angular_velocity-driven stick rotation), retuned damping/friction/density for less twitchy contact. - zarr_writer: bump episode metadata version 0.1 -> 0.3. - 7 new smoke tests in tests/test_features.py (all 34 tests pass). Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Adds ZarrDataset._read_span for variable-length, no-pad span reads and a new ZarrEpisodePackedDataset + pack_collate that flatten full episodes (or sequential chunks) into one stream with cu_seqlens. MultiDataModuleWrapper now picks pack_collate per-dataset when the dataset is packed. Includes a pushT smoke script under scripts/. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Replace the recursive monolithic HNet with a flat-list of composable stages (EncoderDecoderStage / ChunkerStage / ComputeStage) wired into a recursive tree at construction time. Each stage owns its own inference state dataclass; chunkers register ratio-loss aux on a shared HNetContext. CondEncoderModule is extracted from HNetPolicy as a standalone Hydra-instantiated module. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Packed-mode plumbing across the full stack so trainHydra can train H-Net
on full pushshapes episodes via cu_seqlens batching.
algo (egomimic/algo/hnet.py)
- Remove legacy data_schematic param; use norm_stats (MultiDataset)
for per-feature normalize/unnormalize and key topology.
- Add HNetPolicy.forward_packed: per-sub-sequence BOS shift,
per-sub-sequence pos_emb indexing, packed cond encoding.
- Add HNet._ar_rollout_packed: per-episode policy.generate at exactly
T_ep steps; results padded back to (B, T_max, action_dim).
- process_batch_for_training detects packed batches via cu_seqlens
and forwards meta keys unchanged.
- forward_training / forward_eval dispatch on _packed.
- HNetPolicy.generate gains optional T arg.
- Wire chunk_stats_from_aux through compute_losses so avg_chunk_len /
boundary_rate surface alongside losses.
- bf16 dtype fixes in forward_packed (bos / pos_emb casts).
- Training recipe opt-in knobs (default OFF): init_weights_range,
lr_multipliers, use_parameter_groups, weight_decay. Algo exposes
parameter_groups(base_lr) for the optimizer hook.
stage tree (egomimic/models/hnet_nets/{context,hnet,stages,routing}.py)
- HNetContext gains cu_seqlens / max_seqlen + .packed property.
- Each stage's forward branches on ctx.packed; ChunkerStage saves and
restores ctx around the inner_stage call with chunked-space
cu_seqlens.
- Add apply_optimization_params, HNet.init_weights (1/sqrt(n_residuals)
on out_proj/fc2), HNet.apply_lr_multiplier, HNet.parameter_groups.
- Add chunk_stats_from_aux helper.
- bf16 dtype fix in DeChunkLayer.step.
data path (egomimic/rldb/{zarr,embodiment})
- ZarrDataset.__getitem__ loops per-frame on JPEG decode when
horizon>1 on an image key (single-shot decode_jpeg crashed on the
array of buffers; bounded resample loop then hung 400s).
- MultiDataset._iter_leaves descends into ZarrEpisodePackedDataset;
populate_from_datasets probes each embodiment exactly once.
- MultiDataset.infer_norm_from_dataset uses pack_collate when the
dataset is packed; sample_frac is now a frame budget.
- ZarrEpisodePackedDataset.set_norm_stats_from added as no-op for
trainHydra parity.
- pushshapes.get_keymap puts horizon on obs keys so the padded path
serves per-frame obs.
pl_model (egomimic/pl_utils/pl_model.py)
- configure_optimizers calls algo.parameter_groups(base_lr) when the
algo exposes it; falls back to flat .parameters().
evaluator (egomimic/eval/eval_hnet.py + configs)
- HNetEvalVideo: masked MSE over valid positions per episode, viz
via pushshapes.viz_gt_preds.
- eval_hnet.yaml + logger/csv.yaml.
Hydra configs
- tsimulation.yaml: switch to ZarrEpisodePackedDataset.from_resolver
with chunking="none", batch_size=8, action_horizon=1024 (hard-coded
in the keymap call since top-level keys hit MultiDataModuleWrapper).
- hnet_pushshapes.yaml: drop data_schematic interpolation; bump
action_horizon to 1024; expose training-recipe knobs (all OFF).
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
- test_hnet_nets.py (57): routing, chunk, dechunk, isotropic, stages (padded + packed), HNet assembly, ratio_loss, chunk_stats, STE, RMSNorm, AdaLN. - test_packed_pipeline.py (9): normalize broadcast on padded vs packed; _iter_leaves descent; multi-frame JPEG decode; end-to-end packed stats collection. - test_training_recipe.py (20): apply_optimization_params, init_weights height-scaled init, apply_lr_multiplier per-stage stamping, parameter_groups (default, with bias/norm WD=0, per-stage groups, AdamW-consumable). Plus algo wiring tests for the opt-in init_weights_range / lr_multipliers / use_parameter_groups / weight_decay kwargs. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
- smoke_packed_dataset.py (updated): ZarrEpisodePackedDataset byte-match vs direct zarr reads; writes per-episode MP4s and actions plot with cu_seqlens verticals. - smoke_packed_norm_stats.py (new): end-to-end norm-stats collection on the packed dataset (catches _iter_leaves, pack_collate, populate_from_datasets bugs). - smoke_packed_training.py (new): direct fwd+bwd through HNet stage tree with a packed batch — bypasses algo, exercises stage plumbing. - smoke_packed_training_e2e.py (new): goes through algo path (process_batch_for_training → forward_training → compute_losses → backward); includes normalize. - smoke_packed_validation.py (new): per-episode AR rollout via HNetEvalVideo.compute_metrics_and_viz; metrics + viz frame per ep. - debug_full_episode_mem.py (new): A40 memory probe — fwd+bwd on longest episode (T=958), both default and forced-all-boundary worst case. Prints peak GiB. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Adds sections covering: - Algo-level packed training: HNetPolicy.forward_packed mechanics (BOS shift per sub-sequence, pos_emb indexing, cond packing), process_batch_for_training packed dispatch, chunker logging wiring. - Validation path (packed + AR rollout): _ar_rollout_packed, HNetPolicy.generate(T=T_ep), unpack-to-padded contract. - Eval class HNetEvalVideo + eval_hnet.yaml: masked MSE + viz. - Bug fixes uncovered during wiring (JPEG decode, _iter_leaves, populate redundancy, infer_norm pack_collate / frame budget, no-op set_norm_stats_from, bf16/fp32 mismatches, key_name typo). - Training recipe (init + LR + WD): apply_optimization_params, init_weights with 1/sqrt(n_residuals), apply_lr_multiplier, parameter_groups. Opt-in via HNet algo kwargs (default off). - trainHydra invocation (debug + caveats). - Smoke scripts table. - Tests table (86 tests, ~20s). Also updates the "not implemented" notes to reflect what's now wired. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Four new variants for ablating the H-Net conditioning architecture on
pushshapes packed training:
1. **Delta actions** (data-side):
- New ``DeltaAction`` transform that converts absolute action targets
to per-step deltas (``out[t] = arr[t] - arr[t-1]``, ``out[0] = 0``).
- New ``data/tsimulation_delta.yaml`` wiring the transform onto the
packed pushshapes dataset.
2. **Cross-attention** (block-side):
- New ``CrossMultiHeadAttention`` in ``blocks.py`` — queries from x,
keys/values from cond tokens. Padded + packed forward + step (with
its own KV cache that accumulates cond tokens across AR steps).
- ``TransformerBlock.cond_mode`` selects "adaln" (default) or
"cross_attn"; the cross-attn variant inserts an extra cross-attn
layer between self-attn and MLP.
- ``Isotropic`` threads ``cond_mode`` through to each block;
``build_isotropic`` reads it from the spec dict.
- New ``model/hnet_pushshapes_crossattn.yaml``.
3. **Fused tokens** (policy-side):
- New ``FlatFusedPolicy`` + ``HNetFused`` algo in ``algo/hnet.py``.
Bypasses the H-Net stage hierarchy; instead, builds a single causal
transformer stack and feeds an interleaved sequence
``[c_0, BOS, c_1, a_0, c_2, a_1, ...]`` of length 2T. Action
predictions are extracted from odd positions.
- Supports forward (padded + packed) and AR generate.
- New ``model/hnet_pushshapes_fused.yaml``.
4. **Bigger AdaLN** (capacity-only):
- New ``model/hnet_pushshapes_big.yaml``: same baseline arch but
d_model 128->256, d_intermediate doubled, num_heads 4->8.
All implementations pass CPU instantiation + forward/generate/packed
smoke tests; the existing 57 hnet_nets unit tests still pass.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
**Eval change:** ``HNet.forward_eval`` now runs a single teacher-forced forward pass (via ``forward_packed`` -> unpack to padded) instead of the per-episode AR rollout. AR rollout was meaningless for this eval setup: the obs sequence is fixed (recorded), so the predicted action doesn't change what the model sees next — AR just compounded exposure-bias error without testing anything useful. The previous ``_ar_rollout_packed`` method is preserved (dead code for now) for potential closed-loop sim eval later. **Two new opt-in variants:** 1. ``model/hnet_pushshapes_recipe.yaml`` — baseline AdaLN + the H-Net paper's 3 training tricks (residual-stream-aware init, per-stage LR multipliers [3.0, 1.7, 0.9], parameter_groups with WD=0 on bias/norm + WD=0.01 elsewhere) + LR 5e-5 (half of baseline). 2. ``model/hnet_pushshapes_fused_lowlr.yaml`` — fused tokens (current best) with LR 5e-5 to test whether ep 49->99 val drift was overfitting from too-aggressive LR. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
When ``self.model.parameter_groups(base_lr)`` returns a ``list[dict]``, passing it as a kwarg through ``hydra.utils.instantiate`` wraps the dicts as OmegaConf ``DictConfig`` objects, which trips AdamW's tensor-type validation: TypeError: optimizer can only optimize Tensors, but one of the params is omegaconf.dictconfig.DictConfig Fix: instantiate the optimizer's ``_partial_`` first, then call it with the native-Python params arg. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
**HNetSimEval** (egomimic/eval/eval_hnet_sim.py): - New ``EvalVideo`` subclass wrapping ``Tsimulation.pushshapes. PushShapesEnv`` for closed-loop validation. For each val episode it resets the env, steps it one frame at a time using the policy's predicted actions, and reports mean final ``_coverage()`` plus a binary success rate at a configurable threshold. - Handles both ``HNetPolicy`` (stage-based H-Net, ``policy.hnet.step``) and ``FlatFusedPolicy`` (flat transformer, two-token interleaved cond+action step path). - Init modes (yaml-configurable): ``replay`` (default, reads frame-0 state from batch), ``random`` (env.reset(seed=ep_idx)), ``seeds`` (cycle a fixed seed list). **eval_hnet_sim.yaml**: env / init / metric knobs. **scripts/smoke_sim_eval.py**: standalone smoke that loads a Lightning checkpoint, builds the val dataloader, and runs HNetSimEval on a few episodes. Prints metrics and saves an .mp4. No Lightning dependency. **scripts/plot_variants.py**: multi-panel matplotlib comparison of all variant runs (train loss, emb action loss, val MSE, boundary rate, avg chunk len, ratio loss). Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
…t_step The closed-loop env-rollout logic was previously HNet-specific. Refactored so each algo owns its own inference state + step semantics and the eval class is fully algo-agnostic. **New: egomimic/eval/eval_sim.py — SimRolloutEval**: - Owns env reset/step/render, replay-vs-random init, coverage/success metrics, video buffering. - Per-embodiment obs formatter table (``_OBS_FORMATTERS``) handles the env -> model obs shape conversion. - Calls ``algo.sim_init_state(...)`` once per episode, then ``algo.sim_predict_step(state, obs_norm, t, emb_id)`` per step. **New: HNet.sim_init_state / HNet.sim_predict_step** (egomimic/algo/hnet.py): - ``sim_init_state``: allocates the AR inference cache, primes the initial input token. Handles both HNetPolicy and FlatFusedPolicy. - ``sim_predict_step``: encodes single-frame obs, runs one transformer step, returns normalized action ``(B, 1, action_dim)``. Mutates state. These two methods are the **shared inference API** for AR-style closed-loop rollout. ``robot/rollout.py`` and other consumers can call them directly without going through SimRolloutEval — the eval class is just one consumer. **eval_hnet_sim.py** is now a back-compat shim that re-exports ``SimRolloutEval`` as ``HNetSimEval``. HPT's equivalents (chunk-based ``sim_predict_step``) land in a follow-up commit alongside the new HPT model config + training run. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
…licy - New ``model/hpt_pushshapes_circle.yaml`` + ``data/tsimulation_hpt.yaml`` for current-API HPT training on the circle dataset. - ``HPT.sim_init_state`` / ``HPT.sim_predict_step`` for chunk-based AR rollout (refills via ``policy.forward`` every ``action_horizon`` steps). - Fix ``hpt_nets.ResNet.forward`` to use ``.reshape`` instead of ``.view``. - ``HNetPolicy.init_step_state`` / ``step`` and ``FlatFusedPolicy`` equivalents own AR step semantics; algo methods are now 2-line wrappers. - ``eval_sim``: ``ac_key`` lookup is polymorphic across algos. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Previously ``HNetEvalVideo`` passed only ``unpacked[:, 0]`` to ``pushshapes.viz_gt_preds`` so the validation .mp4 was one frame per episode (trajectory overlay drawn on frame 0). Now passes the full per-episode video and updates ``viz_gt_preds`` to: - Accept 5-D ``(B, T, C, H, W)`` images (per-frame) and the original 4-D ``(B, C, H, W)`` legacy single-frame path. - Read ``seq_lens`` from the batch dict to mask zero-padded tail frames. - Insert a 5-frame black separator between episodes for readability. ``_denormalize_imagenet`` generalised to N-D via ``moveaxis``. New ``scripts/smoke_teacher_eval.py``: loads a ckpt, runs ``HNetEvalVideo.compute_metrics_and_viz`` on the first val batch, writes the per-frame .mp4. Verified on ``fused_lowlr_v3`` (3269 frames @ 384²). Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
**Bug fix**: ``SimRolloutEval._init_env`` was reading state from a batch that had already been through ``process_batch_for_training`` (which normalizes via ``norm_stats``). ``env.set_state`` expects world coordinates (0-512 pixel range) so replay-mode init was placing pusher/object at normalized ~[0, 1] positions, producing nonsense rollouts. Now unnormalize via ``self.model.norm_stats.unnormalize`` before calling ``set_state``. Threads ``emb_id`` into ``_init_env``. **New: scripts/replay_episode_in_sim.py** — confirms the dataset is correct by replaying a recorded episode through the env (raw zarr state + actions, no model). Saves a side-by-side .mp4 (left = recorded JPEG, right = env render). Verified: state matches exactly through t=5 then drifts <15px after t=10 (sim non-determinism). Final coverage ~61% over 200 steps, confirming env physics tracks recorded trajectory. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
When reusing precomputed norm stats across datasets (HPT's per-frame loader on the pushshapes circle data uses stats originally computed on the packed loader), ~1-2% of frames have ``state_agent_obj`` values just outside the quantile_1/quantile_99 range used as the bounds-check window. Each violation triggers a resample-and-retry, which dominates data-loading wall time (HPT 100-ep run had logged 0 epochs after 1h with this behavior). Added ``skip_bounds_check=true`` kwarg on ``MultiDataset.__init__``: the bounds check still LOGS first-occurrence violations but no longer resamples. Wired into ``tsimulation_hpt.yaml``. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
**EvalList / EvalListSideBySide** (eval_composite.py): - ``EvalList`` runs each sub-eval, merges metrics, lets each save its own video to its own dir. - ``EvalListSideBySide(EvalList)`` additionally concats per-emb image arrays along width into one composite .mp4. **PCATokenEval** (eval_pca_tokens.py): - Hooks ``policy.action_out`` to capture final per-frame tokens. - Fits PCA on all val-batch tokens; per frame renders 2D scatter trail. - Logs ``pca_top2_explained`` variance ratio. **BoundaryStripEval** (eval_boundary_strip.py): - Runs ``policy.forward_packed`` and extracts ``boundary_prob`` from each chunker's ``bpred`` in ``ctx.aux``. - Per frame renders a centered window of coloured squares + yellow current-step marker. **smoke_composite_eval.py**: e2e smoke. Verified on baseline AdaLN ckpt: 3269-frame 384x792 .mp4. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
BoundaryStripEval: - Switched to a continuous matplotlib magma colormap; each timestep now occupies one pixel row by default so the strip shows a smooth gradient and visible step-to-step instability. - ``square_height`` -> ``pixels_per_step``, default window 64 -> 256. - Vectorised per-frame rendering. evaluator/eval_standard.yaml: HNet default. Top-level ``EvalList`` that runs ``EvalListSideBySide(HNetEvalVideo, PCATokenEval, BoundaryStripEval)`` as one composite .mp4 + ``SimRolloutEval`` as a separate sim .mp4. evaluator/eval_hpt_standard.yaml: HPT default, just SimRolloutEval. data/tsimulation_hpt.yaml: batch 32 -> 128, workers 2 -> 8 to fix the per-frame loader bottleneck (HPT v1/v2 made <1 epoch in 60 min). Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
User feedback — fully continuous greyscale is more intuitive for "boundary tick on a timeline": P=1 (boundary fires) = black, P=0 (chunker quiet) = white. Replaces the magma colormap. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
PCATokenEval now hooks the innermost ComputeStage.main_network to capture chunker-output tokens of shape (T_chunked, d_inner) — one token per CHUNK, not per frame. Per video frame t we map to its containing chunk via cumsum(boundary_mask) and plot that chunk's PCA point, so the scatter STAYS PUT inside a chunk and only jumps when a new boundary fires. Falls back to action_out input (one token per frame) for FlatFusedPolicy where there's no chunker. Logs pca_n_tokens. BoundaryStripEval now also captures boundary_mask and overlays red 1-pixel rows in the strip wherever the chunker COMMITTED a boundary, so the strip shows: greyscale background = soft P(boundary), red lines = committed chunk dividers, yellow centre = current frame. Adds the boundary_mask_rate metric. Smoke b8dwmicpd: pca_n_tokens=564 vs ~3234 frames total, boundary_mask rate=0.174 confirms chunker-token capture. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
With pad_h=min the side-by-side composite was cropping every panel to the shortest one — the boundary strip (window*pps = 256px) — chopping ~1/3 off the PCA scatter (384px) and the HNet trajectory video. Switch to pad_h=max so short panels are zero-padded up to the tallest instead. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Always dispatch scripts through an allocated SLURM job; never run on the sky1 login node (no GPU). Documents the salloc + srun --jobid pattern and the squeue lookup for picking an allocation. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
P(boundary) is in [0, 1] and was being used for BOTH the continuous
gradient AND a binary red overlay — confusing. Split them:
A) Gradient strip = continuous greyscale of soft P(boundary).
B) Discrete strip = binary red/white of boundary_mask (committed
chunk dividers).
Stacked side-by-side per chunker, 2-px black divider between (kept even
so the composite stays divisible-by-2 for x264).
Also defensively even-pad the side-by-side composite output in
EvalListSideBySide so a future panel-size tweak can't trip libx264.
Smoke bt80lauru: (3269, 384, 818, 3), encoded clean.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Stand up the project-local CUDA 12.4 toolkit (under ./cuda-12.4 via
micromamba) and build the three optional CUDA exts against torch
2.6.0+cu124's c10 ABI. With them installed the H-Net fast paths kick in:
- flash_attn_varlen_func replaces SDPA in MultiHeadAttention._forward_packed
- mamba_chunk_scan_combined replaces the Python EMA loop in DeChunkLayer
- Mamba2 ('m' / 'M' arch tokens) becomes available
The EMA loop was the documented "main perf cliff" — it's a sequential
``for t in range(M)`` of CUDA launches, so once the chunker's
boundary_rate grows (currently ~17% on trained AdaLN+recipe) M grows to
~500-700 launches per forward. mamba_chunk_scan_combined collapses that
into one parallel scan.
Pinned versions:
- flash-attn==2.7.4.post1 (2.8+ needs torch 2.7's c10::Error ABI)
- mamba_ssm==2.2.4 (2.3+ needs torch 2.7's c10::Warning ABI)
- causal_conv1d==1.5.0.post8 (1.4.0 has a different csrc layout)
- --no-deps + --no-build-isolation TOGETHER (--no-deps is what stops
pip from upgrading torch to 2.12+cu130 mid-install)
Also dtype-guard MultiHeadAttention._forward_packed so fp32 forwards
(smokes, no-autocast inference) fall back to SDPA — flash_attn_varlen_func
only accepts fp16/bf16.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
scripts/install_cuda_kernels.sh runs the whole flow end-to-end on a compute node: drops micromamba in, builds the project-local cu12.4 toolkit, clears stale pip wheel caches, then installs flash-attn 2.7.4 + mamba_ssm 2.2.4 + causal_conv1d 1.5.0.post8 with the --no-deps --no-build-isolation pair (both required — --no-deps stops pip from upgrading torch). build_cuda_exts.sh / recover_torch_cu124.sh remain as composable pieces for the (rare) case where torch gets bumped accidentally. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Env (Tsimulation/pushshapes/env.py): - float64 throughout: action_space, obs, step input, _get_obs - set_state: angle before position (CoG rotation fix) - get_episode_init(): full init state + obstacles + reset_seed + config_hash Collection (Tsimulation/collect/): - zarr_writer: float64 storage for state/actions/goal/reward - zarr_writer: episode_init JSON attr for deterministic replay - mouse_collect: pre-step obs, 2x window scale, float64 actions - scripted_collect: float64 action computation, init_state metadata Replay (Tsimulation/examples/replay_zarr.py): - Use episode_init (float64 + reset_seed) when available - total_frames trim for zero-padding - --all mode with coverage + drift reporting Sim eval (egomimic/eval/eval_sim.py): - float64 for state init and actions - Fixed reset seed for deterministic replay mode - Filter batches to target embodiment only - max_videos support - pushshapes_sim_stick formatter - eval_hnet_sim_cotrain.yaml: dual sim eval (circle + stick) Tests updated: float64 dtype assertions. Verified: 5/5 scripted + 11/11 mouse episodes replay with zero drift. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
…t_slice; in-mem zarr loader
…ttention, token-dropout
…g, temporal-ensemble, fixed-seed eval
… gitignore scratch
RyanPCo
changed the title
RyanPCo
force-pushed
the
ryanco/pusht-freeform
branch
from
July 24, 2026 01:01
79fa4da to
a697656
Compare
- scripted_collect: BucketTracker + per-frame replay drift validation (DEFAULT_REPLAY_DRIFT_THRESHOLD=0.5, shared _replay_step_loop with replay_zarr); rejection reasons surfaced to caller. - collect/balance.py: bucket assignment helpers used to balance scripted collections across coverage buckets. - collect/gympusht_collect.py: gym-pushT collection front-end. - mouse_collect: expanded recording flow with bucket awareness. - zarr_writer: one-chunk-per-array bulk-write path (matches scripts/rechunk_zarr_dataset.py output). - pushshapes/obstacles.py: obstacle-level definitions feeding new collection variants. - pushshapes/env.py, shapes.py, render.py: support circle_small pusher and align rendering with new shape set. - examples/play_random.py: '--pusher circle_small' option. - examples/replay_zarr.py: _replay_step_loop with early-stop drift check. - README + SCHEMA_NOTES: document the compact one-chunk layout. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
RyanPCo
force-pushed
the
ryanco/pusht-freeform
branch
from
July 29, 2026 23:10
a697656 to
aea3a1b
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
4 participants
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.
No description provided.