PaperlayerView exact revision on arXiv

Source-linked AI summary

Beyond Conservatism: Recoverability-Conditioned Exploration for Model-Based Imitation Learning

Xuanlin Chen, Ziyue Wang, Xunlan Zhou, Yuan-yih Shang, Qiang Wu, Shenghua Wan

arXiv:2609.33336v1cs.LGcs.AI

TL;DR

Conservative model-based imitation can leave uncertain dynamics around expert behavior under-sampled, while task-agnostic exploration may spend interaction on dynamics irrelevant to imitation. RECON uses a separate explorer guided by uncertainty conditioned on recoverability estimated through main-policy imagination. Across control, navigation, and manipulation, it improves imitation, robustness, and model accuracy.

  • Problem

    Conservative data collection can leave uncertain dynamics around the expert distribution insufficiently sampled, while generic uncertainty-driven exploration may prioritize task-irrelevant dynamics.

    Source: §1 · pp. 1–2

  • Method

    RECON pairs a conservative main policy with a separate explorer that prioritizes epistemically uncertain states from which multi-step main-policy imagination predicts a return toward expert-compatible behavior.

    Source: §1 · p. 2

  • Results

    Across control, navigation, and manipulation, RECON improves imitation, robustness, and model accuracy, and on DMC and MetaWorld improves interaction efficiency under matched budgets.

    Source: §1 · p. 3, §7 · p. 10

  • Takeaways & Limitations

    RECON directs interaction toward recovery-relevant dynamics around expert behavior, complementing conservative policy optimization by addressing where the model should be improved.

    Source: §7 · p. 10

    The separate explorer adds model parameters and training computation.

    Source: §6 · p. 9

Abstract

from arXiv · show

Model-based imitation learning (MBIL) improves real-environment interaction efficiency by optimizing policies on imagined rollouts from a learned world model. However, the gap between model-induced and real-environment occupancies makes policy learning sensitive to model error. Conservative MBIL mitigates model exploitation during policy optimization, but when real-environment interactions are collected by the same conservative policy, uncertain regions around the expert distribution remain insufficiently sampled. Generic uncertainty-driven exploration, on the other hand, may allocate interaction to novel but task-irrelevant dynamics. We propose REcoverability-CONditioned Exploration for Model-Based Imitation Learning (RECON). RECON separates conservative policy learning from active data collection by maintaining a main policy for task execution and an explorer for real-environment interaction. The explorer is optimized based on epistemic uncertainty conditioned on recoverability estimated from multi-step main-policy imagination, focusing data collection on unknown states from which the main policy can still return toward expert behavior. Experiments on locomotion, navigation and manipulation show consistent gains in interaction efficiency, imitation performance, and robustness, indicating that RECON directs real-environment interaction toward recovery regions around the expert distribution that are underexplored by prior methods, and thereby learns a world model better suited for imitation.

1 INTRODUCTION

Imitation policies can drift from expert demonstrations, while model-based methods trade costly real interaction for imagined rollouts that can amplify model error. RECON separates conservative policy execution from data collection, directing exploration toward uncertain states from which the main policy can return toward expert-compatible behavior.

1 INTRODUCTION

  • Model prediction errors can accumulate in imagined rollouts, and model-induced occupancy can differ from real-environment occupancy.That mismatch makes policy learning sensitive to model error.Source: §1 · p. 2
  • Conservative policy optimization limits model exploitation, but collecting data with that policy can leave uncertain dynamics around expert behavior under-sampled.Action noise, observation error, or unmodeled disturbances can move deployment into states with little model support.Source: §1 · p. 2
  • Task-agnostic uncertainty exploration can spend a limited interaction budget on novel states far from expert behavior that offer little benefit to imitation.Expert demonstrations already provide a task prior for deciding which dynamics matter.Source: §1 · p. 2
  • RECON trains a separate explorer to prioritize states that are both epistemically uncertain and recoverable by the main policy.Multi-step Dreamer imagination and the imitation discriminator estimate whether the main policy can return toward expert-compatible behavior.Source: §1 · p. 2
  • The main policy and explorer share the world model and replay, while only the main policy is deployed.Under matched interaction budgets, RECON improves imitation performance, interaction efficiency, robustness, and coverage of recovery-relevant dynamics on DMC and MetaWorld.Source: §1 · p. 3

2 RELATED WORK

The related work spans interactive and adversarial imitation, conservative world-model optimization, and model-based imitation and exploration. RECON focuses on the data distribution used to learn the world model, conditioning acquisition on recoverability rather than broad novelty.

2.1 REINFORCEMENT AND IMITATION LEARNING

  • Interactive imitation methods gather learner-induced data, while inverse-reinforcement-learning and adversarial approaches infer or optimize a reward-like signal from expert behavior.These methods can correct distribution deviations but typically require substantial real-environment interaction.Source: §2.1 · p. 3

2.2 WORLD MODELS AND CONSERVATIVE MODEL-BASED RL

  • World-model methods can improve sample efficiency but risk model exploitation in regions poorly supported by data.Conservative methods such as MOReL, MOPO, and COMBO account for epistemic model error during policy optimization; RECON addresses the complementary allocation of online data to improve the model.Source: §2.2 · p. 3

2.3 MODEL-BASED IMITATION LEARNING AND ACTIVE MODEL ACQUISITION

  • Prior model-based imitation approaches modify the imitation objective, representation, or controller, whereas RECON studies the world-model training data distribution under limited interaction.Unlike task-agnostic exploration, RECON targets dynamics relevant to imitation and estimates recoverability in imagination without a reverse model or synthetic policy labels.Source: §2.3 · pp. 3–4

3 PROBLEM FORMULATION AND ANALYSIS

The analysis separates model error under real visitation from nominal data collection: conservatism guides model use, but recovery performance depends on whether recovery-relevant states are covered. Recovery-weighted acquisition formalizes a criterion that values uncertainty according to recoverability rather than seeking uniform model accuracy.

3 PROBLEM FORMULATION AND ANALYSIS

  • The formulation defines normalized policy occupancy and one-step transition error for a learned model relative to the true dynamics.For analysis, it assumes a bounded latent task reward under which the expert is optimal; the learner never observes that reward.Source: §3 · p. 4

3.1 THE ACQUISITION BLIND SPOT

  • The standard simulation argument evaluates its model-error term under real visitation rather than imagined visitation.A finite distribution-free transfer factor requires the recovery occupancy to be absolutely continuous with respect to the nominal occupancy with bounded density ratio.Source: §3.1 · p. 4
  • Conservatism controls which model predictions are used, but nominal data collection controls recovery error only where recovery occupancy is covered.Source: §3.1 · p. 4

3.2 RECOVERY-WEIGHTED ACQUISITION

  • At each acquisition round, the method freezes the main policy and an expert-compatibility score to define recoverability for candidate states and transitions.Source: §3.2 · p. 4
  • The acquisition objective weights predictive uncertainty by recoverability over a reference distribution of candidate inputs.In practice, the reference distribution is approximated by replay-seeded imagined transitions visited by the explorer.Source: §3.2 · p. 4
  • The architecture separates pessimistic model use by the conservative main policy from recoverability-gated uncertainty used by the explorer for acquisition.Both policies update one world model, but only the main policy is deployed.Source: Figure 2 · p. 5
  • The recovery-weighted criterion targets the part of predictive risk affecting states judged recoverable, rather than requiring an accurate model everywhere.For a real query, the stated risk reduction is expressed as expected recoverability-weighted information gain; ensemble disagreement approximates information gain in the practical objective.Source: §3.2 · p. 5

4 RECON: RECOVERABILITY-CONDITIONED EXPLORATION

RECON extends a DreamerV2-based adversarial imitation loop with separate conservative and exploratory actors that share model components. The explorer uses imagined recoverability to gate uncertainty for data collection, while deployment uses only the main policy.

4 RECON: RECOVERABILITY-CONDITIONED EXPLORATION

  • RECON adds a separate recoverability-conditioned explorer to the DreamerV2 latent adversarial-learning loop while sharing the representation, world model, discriminator, and replay.Source: §4 · p. 5
  • DreamerV2 represents its latent state with deterministic and Gaussian stochastic components, and trains the model on expert and online replay trajectories using reconstruction and KL terms.Policy rollouts use the recurrent prior, while a bootstrap transition ensemble provides an epistemic score.Source: §4 · pp. 5–6
  • Both actors use the shared discriminator as their imitation signal, but their imagined rewards differ: the main policy uses uncertainty pessimistically and the explorer values only recoverability-conditioned uncertainty.With α = β = 10, the main actor retains CMIL’s pessimistic model-use objective.Source: §4 · p. 6
  • RECON estimates recoverability by rolling the frozen main policy forward from the explorer’s imagined endpoint and scoring whether it regains main-compatible behavior.A single terminal-derived score is shared across the imagined explorer rollout, and the explorer update retains the pathwise gradient through the recovery rollout.Source: §4 · p. 6
  • During each real episode, explorer windows alternate with main-policy steps, and their transitions update shared components; evaluation deploys only the deterministic main policy.A 500-step episode contains 150 explorer steps and 350 main-policy steps.Source: §4 · p. 6

5 EXPERIMENTS

Across eight environments, RECON’s training-time acquisition improves imitation and robustness while directing interaction toward recoverable, expert-near dynamics. Ablations and diagnostics indicate that recoverability conditioning and the design of its gradient path matter to these gains.

5 EXPERIMENTS

  • Across eight environments, RECON achieves the best or competitive imitation performance, with pronounced gains on Walker Run, PointMaze, and Hammer.It also outperforms CMIL on Faucet Close and Handle Press; the methods share the deployed actor and conservative policy objective, isolating acquisition differences.Source: §5 · p. 7, Figure 3 · p. 7
  • RECON improves perturbation performance across evaluated tasks, especially under delay and altered dynamics, when only its main policy is evaluated.Table 1 reports robust deployment of the main policy as CMIL/RECON means over three seeds; DMControl values are returns, while Maze and Meta-World values are success percentages.Source: §5 · p. 7, Table 1 · p. 7
  • In PointMaze, RECON shifts acquisition toward near-support states and away from far-OOD novelty, while CMIL leaves recovery dynamics near the expert trajectory less covered.CMIL visitation is higher along the wall after inertia-induced deviations; Figure 4 compares visitation and acquisition under the same training budget.Source: §5 · p. 8, Figure 4 · p. 8
  • Removing recoverability while retaining the uncertainty bonus preserves early learning but degrades later performance, most clearly on Walker Run and PointMaze.The ablation indicates that uncertainty alone identifies under-specified dynamics but not which uncertainty is useful for imitation.Source: §5 · p. 8, Figure 5 · p. 8
  • The design ablation favors multi-step main-policy imagination over random collection or a terminal-only signal; the 15-window schedule balances coverage and main-policy control.With 10-step collection bursts, retaining gradients through the recovery rollout outperforms stop-gradient, while tail aggregation has comparatively little effect.Source: §5 · pp. 8–9, Figure 6 · p. 9
  • RECON generally reduces world-model prediction error and ensemble uncertainty, and its predicted recoverability is more strongly correlated with empirical expert-tube re-entry than CMIL’s.Figure 7 reports lower held-out error across distance-based coverage regions; replay error may increase when exploratory trajectories are included.Source: §5 · p. 9, Figure 7 · p. 9, Table 2 · p. 9

6 DISCUSSION AND LIMITATIONS

RECON targets recovery-relevant dynamics rather than improving the world model uniformly. Its separate explorer adds model parameters and training computation, while future work could improve acquisition-signal calibration and reliability.

6 DISCUSSION AND LIMITATIONS

  • The separate explorer adds model parameters and training computation; future work could improve uncertainty calibration and recoverability estimates beyond finite-horizon imagination and discriminator scores, especially farther from current data support.Source: §6 · p. 10

7 CONCLUSION

RECON complements conservative model use by directing limited interaction toward improving the model. The conclusion reports gains across control, navigation, and manipulation, alongside reproducibility details and formal notation assumptions.

7 CONCLUSION

  • Across control, navigation, and manipulation, RECON improves imitation, robustness, and model accuracy while complementing conservative policy optimization with targeted model improvement.Source: §7 · p. 10
  • The authors report that generative AI assisted writing and literature retrieval; they reviewed AI-assisted content and checked AI-identified literature against original sources.Source: AI USE STATEMENT · p. 10
  • Appendix B specifies training, model and policy hyperparameters, environments, robustness protocol, model evaluation, and baseline implementations; code and implementation are planned for public release upon acceptance.Source: REPRODUCIBILITY STATEMENT · p. 10

A.1 NOTATION AND REGULARITY CONDITIONS

  • The notation assumes 0 < γ < 1, bounded common measurable rewards, and Markov kernels on standard Borel spaces; the acquisition gate and reference distribution remain fixed during query evaluation.Source: §A.1 · p. 13

A.2 PROOF OF THEOREM 1

The proofs relate model error to recovery error through occupancy comparisons, with finite uniform bounds requiring suitable coverage. The second theorem motivates targeted rather than global model accuracy and retains an unavoidable residual term.

A.2 PROOF OF THEOREM 1

  • Theorem 1’s proof averages transition error under real visitation; using learned-model occupancy instead would require a separate change-of-measure argument.Source: §A.2 · p. 13
  • When the recovery occupancy is absolutely continuous with bounded density ratio relative to nominal occupancy, recovery error has a finite uniform comparison factor; without absolute continuity, no finite uniform factor is guaranteed.Source: §A.2 · p. 14

A.3 PROOF OF THEOREM 2

  • Theorem 2’s bound retains a residual term unless the acquisition gate upper-bounds recovery occupancy everywhere, making its motivation targeted rather than global model accuracy.Source: §A.3 · p. 14

A.4 PROOF OF THEOREM 3

The appendix develops the information identity behind the acquisition objective and its cellwise specialization, then describes RECON’s pathwise gate gradient and shared DreamerV2 training setup. It also specifies the evaluation environments and diagnostic labels.

A.4 PROOF OF THEOREM 3

  • With a fixed gate and reference distribution, the proof relates the acquisition-risk reduction to gate-weighted conditional mutual information between prospective observations.The gate and reference distribution must remain fixed during a query evaluation.Source: §A.4 · p. 15
  • Cellwise specialization makes predictive information zero across distinct cells; for deterministic successors, disagreement-based information equals the predictive entropy, but is only a proxy for arbitrary stochastic dynamics.The analysis assumes factorized cell parameters and observations dependent only on their own cell’s parameter.Source: §A.4 · p. 15

A.5 LATENT-MODEL SPECIALIZATION AND PATHWISE GATE GRADIENT

  • The pathwise acquisition gradient combines an epistemic-information term with a term that moves explorer endpoints toward states where the frozen main policy earns a large recovery score.Freezing model, discriminator, and main-policy parameters preserves the input Jacobians needed for the recovery-gradient term.Source: §A.5 · p. 15

B.1 DREAMERV2 BACKBONE AND ACTOR–CRITIC OPTIMIZATION

  • RECON uses a Gaussian-latent DreamerV2 ensemble world model, while the main and explorer share latent optimization machinery but differ in rewards and collected state distributions.Ten bootstrapped prior heads provide predicted means for the epistemic score; expert and online sequences train the shared representation and model.Source: §B.1 · p. 16

B.2 COMPLETE TRAINING ALGORITHM

  • The training loop uses a conservative main policy for task behavior and a separate explorer that collects real transitions using recoverability-gated uncertainty.The shared world model is updated from expert and online data, and the main policy is imagined for discriminator updates.Source: §B.2 · pp. 16–17

B.4 ENVIRONMENTS AND EXPERT OBSERVATIONS

  • The experiments cover locomotion and contact dynamics, bottleneck navigation, and visually similar manipulation tasks with distinct contact objectives; PointMaze support labels are diagnostic only.The support, near-support, and far-OOD labels are unavailable to training.Source: §B.4 · p. 17

B.5 ROBUSTNESS PERTURBATION PROTOCOL

The protocol evaluates methods under matched interaction budgets and perturbed deployment conditions, defines a latent-space open-loop prediction gap, and describes the reproduced baselines. The supplied passages specify the evaluation design, not comparative robustness outcomes.

B.5 ROBUSTNESS PERTURBATION PROTOCOL

  • Robustness evaluation uses three seeds per method–environment pair across clean, noisy, delayed, and dynamics-shifted conditions, with task-specific return or success metrics.Impulses test localized recovery; action changes test actuator mismatch; mass and friction changes test persistent dynamics shifts.Source: §B.5 · pp. 17–19

B.6 OPEN-LOOP MODEL PREDICTION GAP.

  • The open-loop prediction gap compares predicted latent features with posterior features inferred from held-out real observations, averaging discounted H-step errors across valid starting points.World-model rollouts use recorded actions without conditioning on later observations.Source: §B.6 · p. 19

B.7 BASELINE IMPLEMENTATION DETAILS

  • Baselines share expert replay, image resolution, action bounds, wrappers, seeds, and interaction budgets; CMIL is the closest comparison because RECON retains its deployed objective but adds an explorer and recovery gate.V-MAIL is also given the DreamerV2 backbone used by CMIL and RECON to avoid confounding the comparison with world-model differences.Source: §B.7 · p. 19

B.8 COMPUTE, RUNTIME, AND MEMORY OVERHEAD

The appendix describes an artifact-level compute audit and reports that RECON adds a separate explorer to the CMIL setup.

B.8 COMPUTE, RUNTIME, AND MEMORY OVERHEAD

  • RECON shares CMIL’s encoder, ensemble RSSM, decoder, discriminator, and replay, adding one actor, two critics, and their targets; its recovery rollout reuses the shared model and main actor.The cited passage says Table 4 reports a checkpoint and throughput audit, but supplies no overhead values.Source: §B.8 · p. 20, Table 4 · p. 20

C ADDITIONAL EXPERIMENTS

Additional experiments assess whether multi-step main-policy imagination predicts recoverability, whether recoverability conditioning improves data acquisition efficiency, and how disagreement relates to held-out model error. The reported patterns support targeted acquisition without requiring calibrated pointwise error estimates.

C.1 RECOVERABILITY PREDICTION ACROSS IMAGINATION HORIZONS.

  • Longer imagination horizons improve recoverability prediction AUROC and bring the learned discriminator-based gate closer to the oracle on PointMaze.Empirical labels and the oracle use true environment-space distance, while the learned gate uses discriminator scores along imagined main-policy rollouts.Source: §C.1 · p. 20, Figure 9 · p. 20

C.2 ACQUISITION EFFICIENCY.

  • RECON achieves the largest uncertainty reduction on fixed main-policy anchor transitions; ungated uncertainty-only collection is less efficient than reward-only V-MAIL.The reported comparison indicates generic uncertainty seeking can spend interaction on novel dynamics that do little to reduce uncertainty around the main policy.Source: Figure 9 · p. 20, §C.2 · p. 21

C.3 DISAGREEMENT AND HELD-OUT MODEL ERROR.

  • Within expert-support regions, larger ensemble disagreement generally tracks larger held-out prediction error, but the relationship weakens and becomes non-monotone farther from support.Near-support and far-OOD long-horizon error also reflects compounding rollout error and sparse data support, so disagreement serves as a ranking signal rather than a calibrated pointwise estimate.Source: §C.3 · p. 21, Figure 10 · p. 21

Related papers

Loading 2609.33336v1…