influence_refinement

Influence-based training data refinement for robot learning policies, following the CUPID approach [[TODO: CUPID paper citation]]. Each iteration, training episodes are scored by their influence on rollout performance; the highest-scoring episodes are kept, low-scoring ones removed, and the policy is retrained on the curated dataset.

Installation

pixi shell -e influence-refine

lerobot defaults to a local editable checkout at ../lerobot/. If you don't have one, comment that line and uncomment the PyPI fallback in pixi.toml.

fast-jl is declared as a pixi dependency and installs automatically. It contains custom CUDA kernels that need to be compiled — if no pre-compiled binary exists for your CUDA version, the install will try to compile from source and requires nvcc + ninja to be available. If pixi fails on this step, install it manually inside the pixi shell:

pip install --no-build-isolation fast-jl

Quickstart

python experiments/pusht_example.py --config configs/pushT_example.yaml \
  --override policy_ckpt_path=path/to/checkpoint

Policy compatibility

Tested with our DiTFlow policy (lerobot_policy_ditflow). Three functions are specific to flow-matching and assume velocity_net, noise_distribution, and config.do_mask_loss_for_padding on the policy or a direct child module: compute_per_sample_objective (influence_computation/grad_utils.py), patched_mode_level_sampling, and patched_policy_t_end (utils/exploration_sampling.py). Two further functions — grad_embedding and InfluenceComputer.value_features — require the standard lerobot _prepare_global_conditioning interface. To add a new policy, reimplement these functions for your architecture.

The loop

Each iteration (experiments/pusht_example.py, configured via configs/pushT_example.yaml):

1. Collect rollouts — run the policy in the environment to generate evaluation episodes.
2. Collect include episodes — gather exploratory rollouts (action-level and/or mode-level) that serve as the influence query set.
3. Compute influence — score each training episode by its influence on the include rollout distribution using TRAK-style projected gradients (Park et al., 2023) with fast CUDA Johnson-Lindenstrauss projections (Fandina et al., 2022).
4. Select — keep the top-k episodes for the include set; optionally remove the lowest-scoring episodes from the training set.
5. Retrain — fine-tune on the merged selected dataset.
6. Repeat with the updated policy.

Anchor-aware training sampler

When meta/segment_anchors.json is present in the training dataset root, sampling during lerobot training is restricted to the annotated frame ranges. This is handled transparently by utils/policy_utils.py via a monkey-patch of lerobot's EpisodeAwareSampler — no changes to lerobot are required and any policy architecture is supported.