Learning on-Policy Learning to Specialize Robot Policies

Abstract

We introduce Generative On-Policy Learners (GOLs), generative models that quickly adapt policy parameters to an unknown target domain. Conditioned on an initial (potentially random) policy and a few interaction episodes in the target domain, GOLs learn to generate improved policy parameters, by mimicking the policy evolution in traditional on-policy reinforcement learning. Unlike prior approaches, GOLs do not require (1) knowledge about domain parameters and rewards, (2) a large number of target domain interactions, or (3) expensive online domain identification. Generation is performed in a behavior-aligned latent space. We evaluate GOLs on continuous control problems, where they outperform state-of-the-art baselines in few-shot domain adaptation on the most dynamic tasks. In addition, we demonstrate that our weight generation approach can iteratively improve policies in a sim-to-real transfer task. Finally, we characterize the key properties of the training dataset required to effectively train the weight generator, offering insights for scaling generative policy synthesis.

Method

GOL generates the parameters of an MLP policy with a desired return, given the parameters of an initial policy and a single state-action trajectory of that policy in the target domain. It uses a trajectory-conditioned auto-encoder, pre-trained with a DAgger-inspired behavior-alignment loss, and a reward-conditioned latent propagator.

Real-world sim-to-real adaptation

We adapt simulation-trained specialist policies to a real Crazyflie 2.1 Brushless quadrotor on a lemniscate track, with no system identification. Each adaptation step uses at most 6 s of flight. Initial policies fail to complete the track; after two steps all of them fly reliably. The video at the top of this page follows one policy from its initial crash through two adaptation steps.

Left: lap times per adaptation step, compared to RAPTOR* and DR-H. Right: flight paths for two policies across iterations.

Simulation: few-shot adaptation

Policies pretrained in source domains are adapted to unseen target domains (leg rescaling for the ant and legged robot, per-propeller thrust loss for drone racing) using one trajectory per iteration and no target-domain rewards.

Average episode reward of policies adapted to unseen in-distribution (ID) and out-of-distribution (OOD) target domains. Results are averaged over 4 training seeds; for each seed, GOL results are averaged over 16 initial policies from held-out test runs.

Ablations

Ablations on drone racing (mean episode reward, single training seed).

w/ BC loss:

the auto-encoder is trained on states visited by the input policy instead of the generated policy.

w/o input weights:

the generator sees only the trajectory and the target reward.

w/o intermediate checkpoints:

training targets come only from the top 10% of checkpoints in each run.

w/o reward normalization:

raw returns are used as the reward condition.

Intermediate checkpoints as training targets and the input weights are the most important design choices, followed by the behavior-aligned loss. Raw returns perform comparably ID and better OOD, but require the specialist's return in the target domain, which normalization avoids.

Sensitivity to the RL algorithm used for training data

GOL adapts much better when trained on SAC checkpoints than on PPO checkpoints (default ant environment). SAC policies have a lower effective rank throughout training and, after the first 20% of training, smaller gradient magnitudes, which indicates a smoother optimization landscape. We hypothesize this favors weight generation, since small prediction errors then degrade performance less.