Source-linked AI summary
Asymmetric On-Policy Distillation Learns In-Context Exploration
Sriyash Poddar, Yanda Bao, Jacob Krantz, Matthew Chang, Xavier Puig, Roozbeh Mottaghi, Natasha Jaques, Abhishek Gupta
TL;DR
Agents in partially observed tasks must explore to infer hidden information, but ordinary imitation and offline context coverage do not provide a scalable route to adaptive exploration. ASTEROID trains a history-conditioned student on its on-policy contexts with privileged-expert labels. Across 11 environments, it outperforms RL and supervised baselines under fixed exploration budgets while using two orders of magnitude fewer training interactions.
Problem
Method
Results
Across 11 environments, ASTEROID outperforms RL and supervised baselines under fixed exploration budgets while using two orders of magnitude fewer training interactions.
Source: §1 · p. 3Takeaways & Limitations
The student can search while uncertain and exploit once the hidden state is revealed, approximating Bayesian posterior sampling on covered histories.
Source: §1 · p. 2, §5 · p. 6The student cannot reduce uncertainty when hidden information can only be gathered outside expert support, and imitation may produce unsafe exploration.
Source: §7 · p. 14
Abstract
from arXiv · showhide
Sequential decision-making under partial observability is difficult because agents must explore to infer hidden, task-relevant information before acting effectively. Learning such information-gathering behaviors tabula rasa via online reinforcement learning is sample-inefficient, whereas offline imitation averages expert actions across hidden states and fails to learn exploration. In this work, we propose ASTEROID, a simple and effective student--teacher supervised learning framework that can learn test-time exploration. Although a privileged expert never needs to explore, we demonstrate that a history-conditioned student can learn to do so when trained with on-policy contexts: the student first collects on-policy context under partial observability, and the privileged expert provides action labels conditioned on that context. By distilling expert actions directly conditioned on the student's on-policy history, the student learns to systematically resolve partial observability, i.e., search while uncertain and exploit once the hidden state is revealed. This enables efficient test-time exploration (similar to Bayesian posterior sampling) without incurring the impractical training-data costs of reinforcement learning. Across a diverse range of eleven different environments, including discrete exploration benchmarks, robotic manipulation, navigation, and Procgen environments, ASTEROID achieves ~2x higher returns under fixed exploration budgets, uses ~100x fewer interactions than RL baselines, and outperforms frontier models at exploration. Finally, our method transfers policies trained in simulation to vision-denied real-world robot manipulation, exhibiting efficient exploration and increasing success by +40% over standard distillation. Overall, it provides a simple yet surprisingly effective framework to learn efficient test-time exploration.
1 Introduction
Agents in partially observed tasks must gather hidden information before acting, but online RL is sample-inefficient and ordinary expert imitation does not teach adaptive exploration. ASTEROID trains a history-conditioned student on its own on-policy histories with privileged-expert labels, enabling it to explore while uncertain and exploit once the hidden state is revealed.
1 Introduction
- Sequential decisions under partial observability require agents to explore for hidden, task-relevant information, as in a household robot searching for an object.The agent must track locations already searched and adapt future exploration based on prior attempts.Source: §1 · p. 1
- Online reinforcement learning is sample-inefficient, while averaging expert actions across possible hidden states can make a student dither rather than resolve uncertainty.Privileged experts act with hidden-state access and therefore do not demonstrate how to find that state.Source: §1 · pp. 1–2
- ASTEROID trains a history-conditioned student on its own partial-observation histories and uses a privileged expert to label the actions appropriate to the states reached.The expert does not explore; failed searches and probes become context for the student’s next decisions.Source: §1 · p. 2
- On-policy contexts let the student learn to rule out failed hypotheses and explore remaining possibilities, then exploit once the hidden state is revealed.The paper describes this behavior as in-context exploration and reports that a curriculum over rollout horizon improves sample efficiency and longer-context adaptation.Source: §1 · p. 2
- ASTEROID approximates Bayesian posterior sampling empirically and theoretically without maintaining an explicit posterior or requiring impractical offline data costs.The paper identifies explicit-posterior methods and infinite offline-data assumptions as barriers to scaling posterior sampling.Source: §1 · p. 3
- Across 11 environments, ASTEROID outperforms RL and supervised baselines under fixed exploration budgets while using two orders of magnitude fewer training interactions.Evaluations include realistic household navigation and real-world robot manipulation.Source: §1 · p. 3
2 Related Work
Prior work addresses exploration through interactive imitation, meta-RL, posterior sampling, and in-context learning, but these approaches face supervision, interaction, or data-coverage constraints. ASTEROID instead distills expert actions on student-generated histories.
Student–Teacher Distillation
- Interactive imitation addresses covariate shift, but prior exploration approaches have relied on extra assumptions, hidden-information prediction, or RL bootstrapping.ASTEROID instead iteratively distills expert actions conditioned on on-policy histories.Source: §2 · p. 3
Meta Learning
- Meta-RL learns adaptive exploration across task distributions but often requires extensive online interaction, while ASTEROID uses supervised learning for long-horizon exploration.ASTEROID also targets POMDPs where exploration occurs within a single episode to estimate the underlying state.Source: §2 · p. 3
Posterior Sampling
- Posterior-sampling methods maintain explicit beliefs and solve for policies under sampled parameters, limiting scalability in high-dimensional settings.ASTEROID does not optimize a posterior-sampling objective; posterior-sampling-like behavior emerges through iterative on-policy distillation.Source: §2 · p. 3
In-Context Decision-Making
- DPT pairs randomly generated contexts with privileged supervision but assumes offline data covers all possible histories, which is intractable in high-dimensional domains.ASTEROID uses on-policy contexts to remove that coverage assumption.Source: §2 · p. 3
3 Problem Statement
The paper studies decision-making in POMDPs, where policies use observations and history rather than the hidden state. It also defines BAMDPs as a special case with a fixed latent parameter.
3 Problem Statement
- A POMDP is represented by state, action, observation, emission, transition, reward, initial-state, and discount components, while the policy receives partial observations rather than the state.The policy’s history includes earlier observations, actions, and rewards.Source: §3 · p. 4
- A BAMDP is a POMDP whose only hidden variable is a latent parameter z, sampled from a prior and held fixed across K episodes.Given z, the BAMDP is fully observed and can be represented as a collection of MDP tasks.Source: §3 · p. 4
- In the three-door example, the agent must use interaction history to infer the hidden target and then navigate to it.The target is sampled uniformly at the start of an interaction and remains fixed across multiple episodes.Source: §3 · p. 4
- Training uses a privileged teacher with access to the true state, while the deployed student acts from partial observations and its history.The learner collects transitions in the POMDP and queries the teacher for action labels.Source: §3 · p. 4
4 Background: Student–Teacher Distillation
Under partial observability, naive imitation of privileged experts does not teach adaptive exploration, and DPT’s offline context coverage requirement is impractical at scale. DPT can learn exploration under coverage assumptions, whereas ASTEROID motivates learning from on-policy student histories.
4 Background: Student–Teacher Distillation
- Naive behavior cloning of privileged trajectories teaches marginal expert behavior rather than adaptive exploration when the student cannot distinguish hidden goals.In the three-door example, identical initial observations lead to learning a uniformly random door choice.Source: §4 · p. 4
- Naive asymmetric distillation learns i.i.d. exploration but fails to adaptively exploit sequences of rewards and feedback.A privileged teacher solves the task without exploration, unlike the partially observed student that must gather information first.Source: §4 · p. 4
- DPT trains a causal transformer from randomly sampled histories, query observations, and privileged expert actions, using the history as context for predicting actions.Under coverage assumptions, it forms an implicit posterior from history and averages expert actions under that posterior.Source: §4 · p. 5
- DPT’s full-coverage assumption requires a dataset to cover the combinatorial space of arbitrary contexts and query observations, which scales exponentially with interaction horizon.Finite-data coverage quickly becomes impractical in high-dimensional observation and action spaces with sparse informative feedback.Source: §4 · p. 5
5 Asymmetric Student Teacher Roll-Out In-context Distillation
ASTEROID replaces exhaustive history coverage with student-generated on-policy contexts labeled by a privileged expert. The section describes its iterative training, deployment, and extension from fixed hidden parameters across episodes to arbitrary POMDPs.
5 Asymmetric Student Teacher Roll-Out In-context Distillation
- On-policy student histories paired with privileged expert labels provide coverage for learning adaptive exploration without covering the entire combinatorial history space.This relaxes the infinite-history-coverage assumption that makes DPT intractable.Source: §5 · p. 5
- Under compliance, on-policy coverage, and bounded episode cost, ASTEROID's regret is posterior-sampling regret plus O(ϵK).The gap to posterior sampling is O(ϵK); an offline in-context learner with the same training error can incur O(ϵK^2) error as test histories drift from its fixed training distribution.Source: §5 · p. 6
5.1 How does ASTEROID learn in-context exploration?
- Each iteration collects student context, then expert trajectories under the same hidden parameters, using expert observations and rewards as queries and expert actions as labels.Repeated iterations expand the available context: after two failed doors, labels point to the remaining door, allowing the student to identify the goal.Source: Figure 3 · p. 6, §5.1 · p. 7
5.2 Practical Implementation
- At deployment, the policy starts with an empty context and builds it through online interaction; a context-length curriculum expands training rollouts across iterations.For evaluations longer than the training context length, deployment uses a sliding window.Source: §5.2 · p. 7
5.3 How does ASTEROID work on arbitrary POMDPs?
- For arbitrary POMDPs, ASTEROID samples a timestep, collects the student's history, and queries the expert from the reached state for action labels and a partial trajectory.The hidden state can change within an episode, so training covers on-policy histories across timesteps; the sampled timestep is biased toward longer horizons to induce a curriculum.Source: §5.3 · p. 8
6 Experiments
The experiments test ASTEROID across partially observed tasks, comparing deployment performance, training efficiency, exploratory behavior, real-robot transfer, frontier models, and on-policy context collection. Across these evaluations, the reported findings favor ASTEROID for efficient exploration under limited sensing.
6.1 Training and Evaluation Details; 6.2 Tasks; 6.3 Baselines
- Evaluation spans partially observed gridworlds, navigation, Procgen mazes, household object search, simulated manipulation, and real-robot manipulation.Methods are evaluated across multiple test-time episodes on held-out goals with a fixed test-time budget; comparisons include ten baselines from several frameworks.Source: §6.1 · p. 8, §6.2 · p. 9, §6.3 · p. 9
6.4 Does ASTEROID learn to solve partially observed problems efficiently on deployment?
- ASTEROID achieves higher final performance than baselines under the same test-time budget, while Procgen Maze shows higher success and 2× faster solving.The authors attribute this to using history to guide search toward unexplored regions rather than repeating searches.Source: §6.4 · pp. 9–10, Figure 6 · p. 10, Figure 7 · p. 10
6.5 Is ASTEROID sample-efficient?
- ∼100× greater environment-step efficiency than RL baselines lets ASTEROID learn robot manipulation policies that explore and successfully pick or insert objects.With history of actions and tactile sensing but no proprioception, ASTEROID reaches ∼92% success in the reported manipulation task.Source: §6.5 · p. 10, Figure 8 · p. 10
6.6 What kind of in-context exploratory behavior do policies learned via ASTEROID display?
- In Tactile-Pick, ASTEROID probes multiple grasp points and uses tactile feedback from failed attempts to guide subsequent actions.In wrist-camera Peg-Insertion, it scans the tabletop to localize the object and target, while naive distillation goes out of distribution.Source: §6.6 · p. 11, Figure 9 · p. 11, Figure 10 · p. 11
6.7 Does ASTEROID learn exploration policies that can be deployed on real robot hardware?
- On the real robot, the student must infer cube location using arm and gripper proprioception alone, then try different grasp locations.The figure reports higher success across 10 trials with randomized cube positions for ASTEROID than naive behavior cloning.Source: §6.7 · pp. 11–12, Figure 12 · p. 12
6.8 Do frontier models perform in-context exploration?
- In wrist-camera Peg-Insertion, ASTEROID achieves ∼78% success versus ∼50% for Claude Opus 5.5; the figure also reports faster exploration.Under full observability, Claude Opus 5.5 reaches ∼65% success, below the expert’s ∼90%.Source: §6.8 · pp. 12–13, Figure 13 · p. 12
6.9 Is learning from on-policy contexts necessary?
- AAWR trained on offline datasets is highly sensitive to dataset quality and diversity, whereas ASTEROID generates student on-policy contexts and is robust to imperfect expert labels.Across domains, DPT performs poorly when offline random contexts fail to capture deployment histories.Source: §6.9 · p. 13
7 Discussion
ASTEROID trains history-conditioned students from privileged experts using student-generated on-policy histories. The discussion situates this approach among prior methods, reports additional POMDP results, and identifies limits tied to expert support and safety.
7 Discussion
- ASTEROID’s iterative on-policy distillation learns in-context exploration across simulated and real-world tasks without reinforcement learning or latent regression.The student collects histories under partial observability and receives privileged-expert targets conditioned on those histories.Source: §7 · p. 13
- The approach requires a strong privileged expert, and exploration cannot reduce uncertainty about hidden information outside the expert’s overall support.The authors also flag potentially unsafe exploration and propose explicit constraints or safety filters as future directions.Source: §7 · pp. 13–14
A Related Work
- Related work contrasts ASTEROID’s on-policy histories with offline-context coverage assumptions, RL-based meta-learning, explicit-posterior methods, and per-step imitation.The paper frames its scope as POMDP exploration within a single episode, rather than only Bayes-Adaptive MDPs with fixed parameters across episodes.Source: A Related Work · p. 20
B.1 Performance of ASTEROID on additional environments
- ASTEROID achieves ∼92% success on Tactile-Pick without proprioception, using tactile forces and the last action to address a high-dimensional POMDP.History-conditioned behavior cloning reaches ∼48%, partly because some evaluation resets need little directed exploration.Source: §B.1 · p. 21
- The additional environments include maze navigation, robot manipulation, and Habitat object search under limited views, random resets, or partial sensing.Tactile-Pick ablations report no performance drop with a noisy expert and improved efficiency with an increasing exploration-horizon curriculum.Source: Figure 16 · p. 21, Figure 17 · p. 21
B.2 Robustness of ASTEROID to noise in privileged expert supervision
The appendix tests noisy supervision, exploration-horizon curricula, scalability, iterative context accumulation, shifted priors, and navigation. It reports robustness to the tested noise and curriculum settings, while showing gradual degradation under prior shift.
B.2 Robustness of ASTEROID to noise in privileged expert supervision
- With Gaussian expert-action noise σ = 0.2, ASTEROID remains performant, with no observed performance drop.The slight improvement may stem from regularization by noisy targets.Source: §B.2 · p. 21
B.3 How does the curriculum of exploration affect ASTEROID during training?
- An increasing exploration-horizon curriculum is more efficient than no curriculum, while performance is relatively insensitive to the curriculum parameter.The curriculum keeps rollouts in distribution and makes on-policy data generation more sample-efficient.Source: §B.3 · p. 22
B.4 ASTEROID is significantly more scalable than RL in complex domains
- ASTEROID performs better than RL baselines in high-dimensional, long-horizon settings, although RL learns competitive search policies in small 5 × 5 Darkroom.The comparison attributes the advantage to supervised expert-action modeling with exploration context built from student rollouts.Source: §B.4 · p. 22, Figure 18 · p. 22
B.5 Does iterative context accumulation actually help ASTEROID?
- In Habitat Object Search, performance improves as exploration steps accumulate and nearly reaches the posterior-sampling upper bound.The reported results also say test-time efficiency improves with additional in-context exploration iterations.Source: Figure 19 · p. 22
B.6 How well does ASTEROID generalize under shifted latent priors?
- In 2D Navigation, performance degrades gradually as evaluation goal radii shift outward from the training radius of 0.7 toward 1.0.Source: Figure 20 · p. 23
C.1 ASTEROID learns efficient exploration in robot manipulation under limited sensing and navigation
- In 2D navigation, ASTEROID adapts its exploration based on context and converges to the true goal.Source: Figure 21 · p. 23
C.2 ASTEROID presents emergent in-context properties such as memory as a result of on-policy in-context learning
ASTEROID exhibits context-dependent memory in Key-Door, while the section contrasts its trajectory-level supervision with DAgger’s per-step labels. The comparison explains how aggregated labels can encourage inefficient exploration.
C.2 ASTEROID presents emergent in-context properties such as memory as a result of on-policy in-context learning
- In Key-Door, ASTEROID uses context to remember a previously found key or door and complete the task without additional exploration after finding the other goal.Source: §C.2 · p. 23, Figure 22 · p. 24
D Per-step interactive imitation learning fails in history-conditioned asymmetric distillation
- DAgger aggregates multimodal per-step labels across hidden goals, leading the student to choose a goal at each step rather than commit to a search hypothesis.ASTEROID instead uses trajectory- or context-level supervision to encourage goal-directed actions over longer horizons.Source: D Per-step interactive imitation… · p. 24, Figure 23 · p. 24
E Didactic Example: ASTEROID learns efficient posterior sampling
The didactic example illustrates how on-policy histories let ASTEROID eliminate inconsistent hidden goals and approach posterior sampling. Training details describe its model, data handling, experts, and compute budgets.
E Didactic Example: ASTEROID learns efficient posterior sampling
- After two failed choices in the three-goal example, ASTEROID’s context identifies the remaining goal, and the student selects it deterministically.Across iterations, the student adapts from uniform choices without context to posterior-sampling-like behavior as on-policy histories accumulate.Source: E Didactic Example: ASTEROID learns… · p. 25
- On-policy contexts more often reach candidate goals and produce feedback that reduces uncertainty, unlike random rollouts that may reveal no reward.Source: E Didactic Example: ASTEROID learns… · pp. 25–26
F.1 Training Details
- ASTEROID uses a causal transformer whose tokens encode observations, previous actions, rewards, and terminal signals; image observations use an ImageNet-pretrained ResNet18.Discrete actions use logits, while continuous actions use a Gaussian mean and standard deviation.Source: §F.1 · p. 26
- For robot manipulation, upweighting new data reduced overfitting to past data, and retaining successful expert trajectories targets actions optimal under parameters consistent with context.Source: §F.1 · p. 26
- Experts are generated by shortest-path planners for navigation and gridworlds and by simulation-trained state-based policies for manipulation.The paper reports a peg-insertion state-based teacher reaching 90+% success in ∼8 h, compared with ∼10% for a wrist-camera policy after 4+ days.Source: §F.1 · p. 26
- Experiments generally use one L40 GPU, and training takes 1–6 h; environment rollouts are the main bottleneck.The appendix says baselines receive as much or more data than ASTEROID, with RL baselines trained to convergence.Source: §F.1 · pp. 26–27, Table 1 · p. 27
F.2 Tasks
The tasks test exploration across robotic manipulation, gridworlds, navigation, indoor object search, and maze navigation under different forms of partial observability. Agents must use interaction histories to reduce uncertainty and act on what they learn.
F.2 Tasks
- Darkroom requires searching for a fixed hidden goal across episodes, while Key-Door adds remembering whether the key was found before searching for the door.Key-Door includes both a version where reward reveals task stage and one where the agent must infer it from history.Source: §F.2 · pp. 27–28
- Ant and 2D Navigation require finding a hidden goal across episodes, using history to avoid revisiting unlikely goals and exploit the goal once discovered.The settings range from continuous 2D navigation to a high-dimensional quadruped, with sparse success rewards.Source: §F.2 · p. 28
- Habitat Object Search requires eliminating checked object locations from a set of candidates and stopping when the onboard detector identifies the target.The privileged oracle planner uses the house map and target location, while the student receives camera observations and detector output.Source: §F.2 · p. 28
- Procgen Maze Navigation requires building an implicit belief from local 3 × 3 observations to avoid uninformative revisits and find a hidden goal.Fixed Maze Navigation holds the maze layout constant to isolate goal-directed exploration from layout generalization.Source: §F.2 · p. 28
- Robot tasks span wrist-camera peg insertion, tactile-only cube pickup, and real-robot cube grasping with proprioception alone.The real-robot evaluation transfers simulation-trained policies to hardware and counts a trial successful when the cube is lifted within 20 seconds.Source: §F.2 · pp. 27–28
G Baseline Implementation Details
The appendix specifies baseline implementations and ASTEROID training details, then analyzes how on-policy histories support posterior-sampling behavior. Its theory relates privileged labels on compliant histories to posterior-sampling supervision and gives an on-support result under stated assumptions.
G Baseline Implementation Details
- Baseline comparisons include DPT, SPOC, AAWR, RL and meta-RL methods, and a Habitat distilled-planner upper bound with access to task structure unavailable to ASTEROID.DPT, SPOC, and AAWR datasets are matched to ASTEROID’s total data; AAWR is sensitive to initial coverage, while the wrist-camera vision RL baseline achieves 10% success after two seeds.Source: G Baseline Implementation Details · pp. 28–29
G.1 Pseudocode
- ASTEROID iteratively collects student on-policy contexts, obtains privileged-teacher labels, aggregates data, and trains the student by conditional imitation.The pseudocode initializes with expert trajectories and expands the sampled context horizon according to a curriculum.Source: §G.1 · p. 29
G.2 Hyperparameters
- The appendix provides separate hyperparameter tables for gridworld tasks and Isaac Lab tasks, with the latter defining an exploration-horizon fraction before expert handoff.The Isaac Lab table describes the exploration horizon as the fraction of an episode during which the exploration policy acts before the expert takes over.Source: Table 3 · p. 30, Table 2 · p. 30
H Additional Theory
- The theoretical analysis frames ASTEROID as DAgger-like learning in history space, contrasting it with offline in-context behavior cloning.The analysis focuses on the BAMDP setting.Source: H Additional Theory · p. 30
H.1 Preliminaries
- The preliminaries model posterior sampling as drawing a hypothesis from the posterior at each episode start and committing to its teacher policy for that episode.The analysis assumes compliant contexts, on-policy support coverage, realizability with exact optimization, and bounded per-episode cost of a wrong hypothesis.Source: §H.1 · p. 31
H.2 Privileged Labels Provide Posterior-Sampling Supervision
- On compliant histories, privileged teacher labels equal posterior-sampling labels, so maximum-likelihood training learns the posterior-predictive policy without explicitly inferring the hidden variable.Under Assumptions 1–3, ASTEROID matches posterior sampling on histories covered by its aggregated on-policy dataset as the context horizon reaches the interaction horizon.Source: §H.2 · pp. 31–32
H.3 Offline versus On-Policy Training
The analysis contrasts offline training on fixed histories with ASTEROID’s aggregation of the student’s own histories. On-policy training removes the history-mismatch term, yielding linear rather than potentially quadratic cumulative error under the stated analysis.
H.3 Offline versus On-Policy Training
- Offline histories can diverge from deployment histories after an earlier error, with a union bound giving mismatch proportional to the number of preceding episodes.The mismatch compounds across episodes even when the offline learner has small error on its training distribution.Source: §H.3 · p. 32
- At least one ASTEROID iterate has average on-policy error at most ϵK + γK, where ϵK is the best in-class error and γK is online regret.The cumulative error bound follows by multiplying the average on-policy supervised error by K.Source: §H.3 · p. 33
- On-policy context aggregation yields cumulative error linear in online episodes, whereas offline history mismatch can make error grow quadratically.The analysis attributes the difference to training on the student’s own histories rather than histories from a fixed behavior distribution.Source: §H.3 · pp. 32–33
- ASTEROID’s Bayesian regret is bounded by posterior sampling’s regret plus a linear approximation term, and finite samples add an eO(K/N) term.The finite-sample statement uses N on-policy contexts per iteration and does not change the linear dependence on the number of episodes.Source: §H.3 · p. 33