cs.ROOct 21, 2026

Asymmetric On-Policy Distillation Learns In-Context Exploration

Authors: Sriyash Poddar, Yanda Bao, Jacob Krantz, Matthew Chang, Xavier Puig, Roozbeh Mottaghi, Natasha Jaques, Abhishek Gupta

Organizations: Paul G. Allen School of Computer Science and Engineering, University of Washington · Work partially done while at Meta

Read PDF
Ask AI
Cite

Copy citation

DOINot available for this paper

Abstract

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.

Figures & tables

Figure 1: ASTEROID provides a simple and reliable method for learning efficient exploration in partially observed domains via on-policy asymmetric student-teacher distillation.
Figure 2
Figure 2: DPT needs intractable full-context coverage for training, while ASTEROID generates on-policy rollouts to learn scalable exploration.
Figure 3: In each iteration, ASTEROID (1) collects on-policy context from the student πθk\pi_{\theta_{k}} and then (2) expert trajectories from π∗\pi^{*} , under the same hidden parameters zz . With this dataset, we can (3) train policies that perform in-context exploration.
Figure 4: We iteratively learn to adapt exploration using feedback over longer contexts.
Figure 5: Tasks: (left to right) meta-learning gridworld benchmark, partially observed navigation, Procgen maze navigation, embodied object search, simulated robot and real robot manipulation.
Figure 6: Performance on discrete exploration tasks requiring memory and goal inference. We observe that ASTEROID outperforms all baselines in terms of test-time efficiency. It performs efficient exploration and almost converges to the expert after inferring the hidden task variables.
Figure 8
Figure 9: End-effector trajectories on the tactile cube-pick task, where the policy must locate and lift a cube using proprioception and touch sensing. The privileged state expert (gray) observes the cube pose and grasps it directly. Under partial observability, the baselines fail to adapt their actions, whereas ASTEROID performs in-context exploration, exploring the table until contact, then grasping and lifting the cube.
Figure 10
Figure 12: (Left) Real robot setup, where it must rely solely on proprioception (arm joint and gripper positions) to infer the object’s location and pick it up. (Right) Comparing the performance of sim-to-real distillation via naive BC and our method on the real-robot task. ASTEROID learns an effective exploration policy that achieves a higher success rate across 10 trials with randomized cube positions.
Figure 13: In the Peg-Insertion task, ASTEROID outperforms a frontier model (Claude Opus 5.5) at exploring the table to pick up the peg and insert it, while also being faster.
Figure 14: ASTEROID generates data online, while AAWR is sensitive to the quality of offline datasets.
Figure 15: We evaluate sim-to-real transfer of policies learned via ASTEROID and naive behavior cloning on a real-robot task, where the robot cannot use the camera to localize the object and pick it up. Instead, it has to use the proprioception of the arm and the gripper and explore different locations by attempting grasps at different positions. We compare the behavior learned with BC vs ASTEROID. The baseline (first row) is trained to imitate expert data only and goes out of distribution. Meanwhile, ASTEROID (second row) learns a policy that uses its in-context exploration to explore multiple grasp locations and pick up the object.
Appendix figures & tables13 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 16: ASTEROID learns to solve POMDPs such as maze navigation, robot manipulation, and Habitat object search, where agents must act under limited fields of view, random resets, and partial sensing.
Figure 17: We run ablations on different parts of ASTEROID in the Tactile-Pick task. Left: we run ASTEROID with a noisy expert and observe no performance drop. Middle: we sweep over different curricula with increasing exploration horizons and find minimal sensitivity. Right: we show that adding curriculum improves the efficiency of ASTEROID by ensuring that as the horizon increases over iterations, the learned policy does not suffer from OOD histories or generating random actions.
Figure 18: RL is effective over short horizons, but ASTEROID uses supervised updates to be effective across long-horizon tasks as well.
Figure 19: The test-time efficiency of ASTEROID increases significantly with more iterations of asymmetric on-policy distillation, demonstrating that as the context dataset grows, the policy can adapt over longer horizons, making exploration more efficient.
Figure 20: Generalization under shifted latent priors in 2D Navigation. Training goals lie at radius 0.7; evaluation goals are sampled at increasing radii up to 1.0. Performance degrades gradually with shift magnitude.
Figure 21: In 2D navigation, ASTEROID adapts its exploration based on context, converging to the true goal.
Figure 22: In the Key-Door task, the agent explores to find both the key and the door to get +1 reward for each task. ASTEROID learns an effective strategy: if it reaches the key or door first, it uses its context to remember the location, and upon finding the other goal, it can immediately complete the task without additional exploration.
Figure 23: Under DAgger, the same student rollout receives per-step expert labels toward different goals depending on the hidden goal (left, middle). Aggregating these labels yields a policy that chooses a goal uniformly at each step (right) instead of committing to one hypothesis.
Figure 24: Toy example demonstrating that ASTEROID learns adaptive in-context behavior, where it effectively explores based on history.
Figure 25: On-policy contexts ensure efficient coverage over relevant histories, while random exploration [ 42 ] can generate uninformative contextual information that harms training.
EnvironmentKK (iters)ASTEROID data (trajs)Offline baselines dataRL baselines env stepsExpert labels (rollouts)Teacher cost
Gridworld (Darkroom, 2D-Nav)1010×1000=1010\times 1000=10 k (env horizon HH )10k expert trajsto convergence10knegligible (planner)
Tactile-Pick33×8192≈243\times 8192\approx 24 k≈\approx 24k expert trajs∼108\sim 10^{8}≈\approx 24k∼\sim 8 h (state-based RL)
Peg-Insertion44×4096≈164\times 4096\approx 16 k≈\approx 16k expert trajs∼108\sim 10^{8}≈\approx 16k∼\sim 8 h (state-based RL)
Habitat55×1024=51205\times 1024=51205120 (Distilled Planner)n/a5120negligible (planner)
Appendix
Table 1: Per-environment data and compute budget for ASTEROID and baselines.
HyperparameterValue
Training
Dataset size (per iteration)1000
ASTEROID iterations10
Batch size64
Learning rate1×10−41\times 10^{-4}
Weight decay1×10−41\times 10^{-4}
Appendix
Table 2: Hyperparameters for ASTEROID on Gridworld tasks.
HyperparameterValue
Training
Demonstrations per iteration8192 for Tactile-Pick / 4096 for Peg-Insertion
Training steps per iteration50,000 gradient steps
Learning rate1×10−41\times 10^{-4}
OptimizerAdamW
Curriculum Schedule (Ratio of episode length)
Appendix
Table 3: Hyperparameters for ASTEROID on Isaac Lab tasks. The exploration horizon denotes the [min⁡,max⁡][\min,\max] fraction of the episode during which the exploration policy acts before handing off to the expert.
More in this topic

RL Exploration

Reinforcement learning exploration studies how agents choose informative actions and states to learn effective policies under uncertainty.

See all 219 papers

Explore similar work

CardsList

  1. Asymmetric On-Policy Distillation: Bridging Exploitation and Imitation at the Token Level

    May 7, 2026Nan Jia, Haojin Yang, Xing Ma +6RL for Language Model ReasoningLanguage Model Distillation

    Huazhong University of Science and Technology · Peking University · Meituan

  2. To Distill or Decide? Understanding the Algorithmic Trade-off in Partially Observable Reinforcement Learning

    Oct 3, 2025Yuda Song, Dhruv Rohatgi, Aarti Singh +1Partially Observable RLLearning Using Privileged Information

    Carnegie Mellon University · Massachusetts Institute of Technology · Aurora Innovation

  3. Explore with Memory, Learn without It: In-Context Exploration and Self-Distillation for LLM Reasoning

    Oct 13, 2026Yutong Chen, Zhichao Xu, Rui Liu +1RL for Language Model ReasoningRL Exploration

    Department of Computer Science, Clemson University · Kahlert School of Computing, University of Utah