[Feature] Integrate RLinf DAgger with EmbodiChain and ODS

[yuecideng]

Background

RLinf now provides an embodied DAgger workflow: the student policy interacts with the environment, the expert relabels the visited observations, the data are aggregated in replay, and the student is optimized with an embodied DAgger loss. RLinf supports synchronous and asynchronous execution, classic in-memory replay, and an online LeRobot rolling-window path.

EmbodiChain already exposes Gym-style environments, explicit observation/action spaces, expert trajectory generation, MotionGenerator/Atomic Action/Task Program execution, and Online Data Streaming (ODS). The current EmbodiChain to RLinf integration is validated primarily for low-dimensional state observations and flat continuous Box actions.

This issue defines the work needed to make RLinf DAgger a first-class workflow for EmbodiChain tasks and to evolve ODS into a reusable learner-in-the-loop data path.

Goals

• Run RLinf DAgger on at least one EmbodiChain task end to end.
• Make observation semantics, action semantics, reset behavior, and timing explicit at the EmbodiChain boundary.
• Support both checkpoint-based policy experts and EmbodiChain-native planner/controller experts.
• Preserve existing expert-only ODS behavior.
• Add an ODS mode that can collect learner-state rollouts, query experts, and aggregate supervision data across iterations.
• Keep the initial implementation focused on vectorized simulation, continuous actions, and synchronous collection.
• Reuse RLinf's DAgger trainer, beta schedule, replay, checkpointing, and distributed runtime instead of duplicating those components in EmbodiChain.

Non-goals for the first milestone

• Reimplementing RLinf's DAgger loss or distributed actor/learner stack inside EmbodiChain.
• Supporting arbitrary Dict, semantic, VLA, or chunked actions without an explicit converter.
• Making Task Program demos state-conditioned before the expert query contract is defined.
• Adding real-robot HG-DAgger in this issue.
• Replacing the current expert-only ODS lifecycle.

Current architectural constraints

• EmbodiChain's native RL trainer is organized around RL rollout kinds and reward/value fields; DAgger is a supervised imitation loop and should not be forced into the PPO/GRPO algorithm registry.
• SyncCollector currently samples a policy action and calls env.step(); it does not query an expert on the same observation.
• EmbodiedEnv expert trajectories are causally aligned but are primarily generated from open-loop demo segments. They are suitable for initial behavior-cloning data, but not for relabeling learner-discovered states.
• ODS currently uses a simulation subprocess to fill a shared-memory buffer and periodically replaces the buffer. DAgger requires aggregation across iterations, retention policy, policy-version metadata, and expert labels for learner-visited states.
• Expert trajectory storage has its own qpos/qvel schema. It must not be assumed to be identical to the policy action space.

Proposed architecture

1. EmbodiChain environment and expert contracts

Add an environment-side expert protocol, owned by the Gym environment layer, with a batched state-conditioned interface:

class ExpertActionProvider(Protocol):
    def reset(self, env_ids, obs, info) -> None: ...
    def act(self, obs, *, env_ids) -> ExpertQuery: ...

@dataclass
class ExpertQuery:
    action: torch.Tensor | TensorDict
    valid: torch.Tensor
    confidence: torch.Tensor | None = None
    metadata: dict[str, object] = field(default_factory=dict)

The contract must specify:

• The observation representation passed to the expert.
• The action representation returned by the expert.
• Whether actions are raw policy actions or controller-ready actions.
• Per-environment validity and optional confidence.
• Reset and terminal-row behavior.
• Device, dtype, batch shape, and finite-value validation.
• Expert failure behavior and safe fallback.

Implement at least:

• PolicyExpert: loads a student/expert checkpoint and performs a batched forward pass.
• PlannerExpert: adapts MotionGenerator, IK, Atomic Action, or another state-conditioned controller.
• An initial-data adapter for existing open-loop expert demos; explicitly mark it as BC-only rather than a DAgger relabeler.

2. Observation and action mapping

Add explicit per-task mapping metadata for RLinf:

• Ordered observation keys and flattening rules.
• Observation normalization and dtype.
• Policy action dimensions and bounds.
• Conversion from policy actions to EmbodiChain raw actions.
• Conversion from planner/expert output to the same raw action space.
• Control frequency, step_dt, action-chunk length, and joint ordering.

Validate these mappings at environment construction and in a worker-local reset/step smoke test.

3. RLinf adapter and first DAgger task

Extend the existing RLinf integration documentation/configuration with a DAgger recipe:

• Use a low-dimensional state task first.
• Use a flat continuous Box action first.
• Run with a checkpoint-based MLP expert first.
• Evaluate only with the student policy.
• Record expert action, executed action, expert-used flag, beta, policy version, episode id, and terminal flags.
• Verify that each expert label corresponds to the learner's pre-action observation.

Then add one planner-expert example after the policy-expert path is stable.

4. ODS DAgger mode

Generalize ODS with explicit producer modes:

• expert_demo: preserve the existing expert-only behavior.
• dagger_sync: learner and expert are queried in the same rollout process.
• dagger_async: learner, environment worker, and replay writer communicate asynchronously.

The DAgger data schema should include:

obs
expert_action
executed_action
expert_used
valid
confidence (optional)
episode_id
iteration
policy_version
beta
terminated
truncated
success

ODS changes should provide:

• Append/aggregate semantics across DAgger iterations.
• Configurable replay capacity and retention strategy.
• Recent-vs-all sampling controls.
• Episode and segment boundary preservation.
• Success-only and failure-focused filters.
• Policy snapshot/version metadata.
• Checkpoint/resume metadata for replay state.
• The existing lock and continuity guarantees for shared-memory sampling.
• Optional export to LeRobot episode format.

The initial async design may use periodic policy snapshots. It must not require a live Python model object to be shared across the simulation process.

5. Native EmbodiChain path (second phase)

Only after the RLinf path works, decide whether EmbodiChain needs a native train-il command. If required, add a separate embodichain/learning/il/ package with:

• DAggerCollector
• DAggerReplayBuffer
• DAggerAlgorithm for behavior-cloning loss
• DAggerTrainer
• Policy action-mean/prediction API
• Evaluation and checkpoint resume

Do not add DAgger to PPO's RolloutKind or make the existing RL trainer pretend that DAgger rollouts contain value/log-probability fields.

Phased implementation plan

Phase 0 — contract and task selection

• Select one low-dimensional EmbodiChain task and one expert source.
• Document observation keys, action semantics, joint ordering, control rate, and reset behavior.
• Decide whether the first expert is a checkpoint policy or a planner.
• Add a worker-local reset/step smoke test.
• Define the minimum metrics and artifact layout.

Phase 1 — EmbodiChain expert and mapping APIs

• Add ExpertActionProvider and ExpertQuery.
• Add shape/device/dtype/finite-value validation.
• Add PolicyExpert.
• Add explicit policy-action to raw-environment-action conversion.
• Add per-environment reset and invalid-label handling.
• Add focused unit tests for batch shape, action semantics, reset rows, and failure masks.
• Update public API docs if the new protocol is exported.

Phase 2 — RLinf DAgger integration

• Add an EmbodiChain DAgger configuration and launch recipe.
• Verify beta scheduling and mixed rollout behavior.
• Verify same-observation expert relabeling.
• Verify student-only evaluation.
• Verify replay sampling, checkpoint save/resume, and policy-version metadata.
• Log expert query latency, invalid-label rate, expert/student action disagreement, beta, success rate, and episode length.
• Add an end-to-end smoke test with a short rollout and one update.

Phase 3 — ODS aggregation mode

• Introduce explicit ODS producer mode selection.
• Add DAgger fields to the shared TensorDict schema.
• Add append/aggregate replay semantics without changing expert_demo.
• Add replay retention and recent/all sampling configuration.
• Add episode commit/finalization and partial-episode discard behavior.
• Add policy snapshot/version handoff for async workers.
• Add tests for concurrent writes, lock windows, continuity, refill, shutdown, and resume.
• Benchmark throughput and memory growth under realistic image/state sizes.

Phase 4 — planner and Task Program experts

• Add PlannerExpert for a MotionGenerator or Atomic Action task.
• Measure planner query latency and determine whether action caching or query decimation is safe.
• Define a state-conditioned Task Program expert API.
• Add planner failure/recovery metadata and valid masks.
• Compare checkpoint-policy expert and planner expert on the same task.

Phase 5 — optional native EmbodiChain DAgger

• Decide whether RLinf already covers the intended workloads.
• If needed, add the separate native IL package and train-il CLI.
• Reuse the same expert and mapping contracts as RLinf.
• Add deterministic evaluation and replay checkpointing.
• Keep RL and IL configuration schemas separate where their lifecycle semantics differ.

Acceptance criteria

• A documented EmbodiChain task runs RLinf DAgger from reset through training and evaluation.
• The student is the only policy used during evaluation.
• Every stored expert label is aligned with the learner's pre-action observation.
• Expert and learner actions pass through the same action conversion and environment preprocessing path.
• Invalid expert outputs are never included in the imitation loss.
• ODS expert-only mode remains backward compatible.
• ODS DAgger mode retains data across at least two aggregation iterations.
• Replay can be resumed from a checkpoint without losing iteration, policy-version, or beta metadata.
• A planner expert can report per-environment failures without corrupting other vectorized environments.
• The end-to-end smoke test completes without leaked workers, unfinished recorder queues, or stale shared-memory rows.
• Throughput, query latency, replay size, and success metrics are reported for the reference task.

Risks and mitigations

• Expert query cost: start with batched policy experts, profile planner queries, and only add query decimation with explicit label-age metadata.
• Action mismatch: enforce one canonical raw action space and test the mapping at construction time.
• Replay memory growth: use bounded replay, recent/all sampling, CPU offload, and optional episode shards.
• Task Program mismatch: keep open-loop demos as BC-only until state-conditioned semantics exist.
• Async policy staleness: attach policy versions to every episode and make snapshot cadence configurable.
• Visual observation size: use disk-backed episode/shard storage or LeRobot rolling windows instead of unbounded GPU replay.

References

• RLinf DAgger: https://rlinf.readthedocs.io/en/latest/rst_source/examples/embodied/dagger.html
• RLinf EmbodiChain integration: https://dexforce.github.io/EmbodiChain/main/guides/rlinf.html
• EmbodiChain ODS: embodichain/data_pipeline/engine/data.py
• EmbodiChain RL trainer/collector: embodichain/learning/rl/utils/trainer.py and embodichain/learning/rl/collector/sync_collector.py
• EmbodiChain expert trajectory path: embodichain/lab/gym/envs/embodied_env.py

Definition of done

The selected reference task has a reproducible RLinf DAgger command, a validated state-conditioned expert path, persistent aggregated labels, documented mapping metadata, passing focused tests, and a measured throughput/memory report.