Source-linked AI summary
Relay Policy Learning: Solving Long-Horizon Tasks via Imitation and Reinforcement Learning
Abhishek Gupta, Vikash Kumar, Corey Lynch, Sergey Levine, Karol Hausman
TL;DR
Long-horizon robotic tasks challenge reinforcement learning and imitation learning because hierarchical methods face practical difficulties and pure imitation can compound errors. RPL first learns hierarchical, goal-conditioned policies from unstructured demonstrations, then fine-tunes them with reinforcement learning; in kitchen simulation, it outperforms hierarchical-RL and imitation-learning baselines. The method is evaluated within a scope that assumes test-time goals come from the same goal distribution as the demonstrations and leaves generalization to longer sequences for future work.
Problem
Long-horizon robotic tasks remain difficult because hierarchical reinforcement learning faces exploration, segmentation, and reward-definition challenges, while pure imitation can be insufficient.
Method
RPL uses relay imitation learning to initialize hierarchical, goal-conditioned policies from unstructured demonstrations, then applies goal-conditioned hierarchical reinforcement-learning fine-tuning.
Results
RPL learns compound long-horizon manipulation tasks and significantly outperforms hierarchical reinforcement-learning and imitation-learning baselines.
Takeaways & Limitations
Unstructured demonstrations can bootstrap a single hierarchical policy that achieves multiple compound goals and remains amenable to reinforcement-learning improvement.
Takeaways & Limitations
The setting assumes test-time goals come from the same goal distribution as demonstration goals, and generalization to longer sequences remains future work.
Abstract
from arXiv · showhide
We present relay policy learning, a method for imitation and reinforcement learning that can solve multi-stage, long-horizon robotic tasks. This general and universally-applicable, two-phase approach consists of an imitation learning stage that produces goal-conditioned hierarchical policies, and a reinforcement learning phase that finetunes these policies for task performance. Our method, while not necessarily perfect at imitation learning, is very amenable to further improvement via environment interaction, allowing it to scale to challenging long-horizon tasks. We simplify the long-horizon policy learning problem by using a novel data-relabeling algorithm for learning goal-conditioned hierarchical policies, where the low-level only acts for a fixed number of steps, regardless of the goal achieved. While we rely on demonstration data to bootstrap policy learning, we do not assume access to demonstrations of every specific tasks that is being solved, and instead leverage unstructured and unsegmented demonstrations of semantically meaningful behaviors that are not only less burdensome to provide, but also can greatly facilitate further improvement using reinforcement learning. We demonstrate the effectiveness of our method on a number of multi-stage, long-horizon manipulation tasks in a challenging kitchen simulation environment. Videos are available at https://relay-policy-learning.github.io/
1 Introduction
RPL addresses the difficulty of applying hierarchical learning to long-horizon robotic tasks by combining unstructured demonstrations with reinforcement-learning fine-tuning. It learns hierarchical, goal-conditioned policies that solve long-horizon manipulation tasks in simulation and outperform the cited hierarchical-RL and imitation-learning baselines.
- Short-horizon robotics successes leave challenging long-horizon tasks as an open problem for reinforcement learning.Hierarchical reinforcement learning introduces temporal abstraction, but practical challenges include exploration, skill segmentation, and reward definition.
- Hierarchical imitation learning typically seeks temporal abstraction and meaningful demonstration segmentation, but is not tailored to subsequent reinforcement-learning fine-tuning.Pure imitation can be insufficient for multi-stage, long-horizon settings.
- RPL uses a two-phase procedure that first pre-trains hierarchical policies from demonstrations and then fine-tunes them with reinforcement learning.The approach is designed to use unsegmented and unstructured demonstrations without requiring explicit skill segmentation or subgoal definitions.
- Fixed-length, goal-conditioned low-level trajectories simplify reward specification and support continuous policy improvement through reinforcement fine-tuning.The low-level policy acts for the same limited length regardless of the goal achieved.
- RPL learns general hierarchical, goal-conditioned policies for long-horizon manipulation tasks in a challenging kitchen simulation environment.The paper reports significant performance gains over hierarchical reinforcement-learning and imitation-learning algorithms.
2 Related Work
Prior hierarchical reinforcement-learning methods provide temporal abstraction but remain challenged by exploration and optimization. Hierarchical imitation-learning methods extract or learn segmented skills, yet traditionally lack support for later reinforcement-learning improvement.
- Hierarchical reinforcement-learning frameworks include options, HAM, max-Q, feudal networks, and multi-layer goal-conditioned policies.The cited methods are presented as promising frameworks for temporally extended tasks.
- Hierarchical reinforcement-learning algorithms still face exploration and optimization challenges.The paper also reports these limitations in its experimental comparison with Nachum et al.
- Hierarchical imitation-learning approaches commonly extract transition segments or learn primitives and latent-conditioned policies from demonstrations.Their objective is to create meaningful segmentation or low-level behaviors.
- Traditional hierarchical imitation-learning methods are generally not designed to improve learned primitives with subsequent reinforcement learning.This limits their fit for challenging multi-task, long-horizon problems where pure imitation may be insufficient.
3 Preliminaries
The paper formulates reinforcement learning as policy optimization in a finite-horizon Markov decision process and extends it to goals. It then defines imitation learning as learning policies from demonstrated state-action trajectories, including goal-conditioned behavior across multiple tasks.
- A finite-horizon Markov decision process is defined by state and action spaces, a transition function, and a reward function.The RL objective is to find a policy maximizing expected reward over policy-induced trajectories.
- Goal-conditioned reinforcement learning learns a policy π(a|s, s_g) that maximizes expected reward with respect to a distribution of goals.The goal-conditioned policy receives both the current state and a goal.
- Imitation learning uses demonstrations containing trajectories of state-action pairs rather than a known reward function.The demonstrations are sequences such as τ_i = {s_i^0, a_i^0, ..., s_i^T, a_i^T}.
- Behavior cloning learns a policy by maximizing the likelihood of demonstrated actions given states.Its objective is max E_(s,a)∼D log π(a|s).
- Goal-conditioned imitation learning extends imitation to demonstrations associated with different goals and learns a policy capable of reaching those goals.The policy is conditioned on the desired goal while imitating the demonstrations.
4 Relay Policy Learning
Relay policy learning combines relay imitation learning from unstructured demonstrations with relay reinforcement fine-tuning. Its hierarchical, goal-conditioned policies use relabeled data at both levels to provide an initialization that can be improved through environment interaction.
- Overview: RPL uses two phases: relay imitation learning initializes hierarchical policies, followed by relay reinforcement fine-tuning for task performance.RIL alone may fail on the most challenging tasks, but supplies an effective initialization for fine-tuning.
- Relay Policy Architecture: The policy architecture combines a high-level goal-setting policy with a low-level subgoal-conditioned policy to generate environment actions with temporal abstraction.The high-level policy selects subgoals every H steps, while the low-level policy acts at every timestep.
- Relay Imitation Learning: RIL learns from unstructured, unlabeled play demonstrations without requiring explicit human goal labels or demonstrations that accomplish every test-time goal.Test-time goals are assumed to come from the same goal distribution as goals accomplished in the demonstrations.
- Relay Imitation Learning: Relay data relabeling creates low- and high-level state-goal-action datasets by treating states reached within fixed windows of demonstrations as reachable goals.The low-level uses window Wl, while the high-level uses Wh and selects subgoal states along demonstration trajectories.
- Relay Reinforcement Fine-tuning: Relay reinforcement fine-tuning improves the initialized policies through goal-conditioned policy gradients while incorporating relabeled demonstrations into the optimization.The method initializes both policy levels with RIL and adds maximum-likelihood terms that keep behavior consistent with the demonstration pools.
5 Experimental Results
Experiments evaluate relay policy learning on compound, multi-stage kitchen manipulation goals, comparing hierarchical imitation and reinforcement fine-tuning against flat policies and learning-from-scratch baselines. RPL improves imitation and fine-tuning performance, while window size and reward choice materially affect results.
- Experimental setup: The evaluation uses a MuJoCo kitchen with compound goals requiring manipulation of multiple scene elements across multiple stages.Tasks include combinations of microwaves, kettles, burners, lights, sliders, cabinets, and hinges.
- Evaluation: Policies are evaluated across 17 compound goals using success and step-completion scores, where each task contains four element-manipulation steps.A step counts when its corresponding element reaches within ϵ of the desired position.
- Relay imitation learning: RIL achieves higher average step-completion scores and fully completed compound goals than flat behavior-cloning variants, although imitation alone is not near-perfect.The comparison includes flat behavior cloning with and without data relabeling.
- Relay reinforcement fine-tuning: RPL fine-tuning significantly improves average success and stepwise completion over baselines, and all three RPL variants outperform flat-policy fine-tuning.RIL initialization at both hierarchical levels also improves performance over HIRO and high-level learning from scratch.
- Relay reinforcement fine-tuning: RPL learns faster and reaches higher success than on-policy HIRO trained from scratch, while hierarchical fine-tuning benefits from easier credit assignment and reward specification than flat fine-tuning.The method also outperforms pre-training only the low level, an oracle segmentation scheme, and a nearest-neighbor open-loop baseline.
- Ablations and analysis: Increasing the low-level window size reduces both imitation and fine-tuning performance because behaviors become more temporally extended.Among tested fine-tuning rewards, sparse reward performs best once exploration is sufficiently directed and is less prone to local optima than the alternatives.
6 Conclusion and Future Work
Relay policy learning solves long-horizon, multi-stage tasks by bootstrapping hierarchical learning from unstructured demonstrations. The authors identify generalization to longer sequences and extrapolation beyond demonstrations as future directions.
- RPL leverages unstructured demonstrations to bootstrap a hierarchical learning procedure for long-horizon, multi-stage tasks.
- A single policy can achieve multiple compound goals requiring temporally extended reasoning.
- Future work targets generalization to longer sequences, extrapolation beyond demonstration data, and off-policy reinforcement learning for improved data efficiency and physical-robot learning.
A Experimental Details
The experiments use feed-forward MLP policies and a standardized on-policy reinforcement-learning setup across compound goals in the kitchen simulation environment.
- All policies use two-layer MLPs with 256 units per layer and ReLU nonlinearities, while flat baselines use the same architecture.Imitation learning uses ADAM with batch size 128 and learning rate 0.005.
- Fine-tuning covers 17 compound goals individually, with path length 280 and low-level horizon 30.The method uses 100 trajectories per on-policy iteration and discount 0.995.
- The reinforcement-learning phase uses a variant of Trust Region Policy Optimization.
- The simulation has a 30-dimensional state space and a 9-dimensional action space with seven arm degrees of freedom and two gripper dimensions.Actions are represented as joint velocities.
B Reward Function Details
The reward-function comparisons evaluate sparse, Euclidean-distance, and element-wise sparse rewards, with sparse rewards used for the reported experiments.
- The reward comparison uses sparse, Euclidean-distance, and element-wise sparse goal-conditioned rewards with ϵ set to 0.3.
- Sparse reward is used for all experimental results in Figure 5.
- Element-wise sparse rewards index distinct scene elements such as the microwave, burners, switches, and cabinets, excluding the robot arm.
C Oracle Baseline Details
The oracle baseline uses hand-designed, semantically meaningful segmentation that produces variable-length demonstration windows. Segment boundaries become high-level actions and goals for low-level goal-conditioned behavior cloning.
- Oracle segmentation: The oracle segments demonstrations when selected kitchen objects move more than ϵ = 0.3, producing variable-sized windows.The monitored objects include the microwave, kettle, light switch, burners, slide cabinet, and hinge cabinet.
- Oracle segmentation: Figure 7 plots different demonstration splits as colored rows over time-steps, showing substantial variability in segment timing.
- Oracle segmentation: Segment-break states serve as high-level policy actions, while low-level policies learn goal-conditioned behavior cloning toward those states.
D Visualization of Learned Behaviors
The paper visualizes learned behaviors for compound tasks, showing successful examples and directing readers to accompanying videos for fuller appreciation.
- The visualizations cover several learned behaviors for compound tasks.
- Successful learned behaviors include opening the microwave, moving the kettle, turning on a light switch, and sliding the slider.
D.1 Successful cases
The visualizations include successful and failed learned behaviors involving kettle movement, controls, sliding, cabinet opening, and the oven light.
- D.1 Successful cases: Successful behaviors include moving the kettle, turning the top knob, sliding the slider, and opening the hinge cabinet.
- D.1 Successful cases: A failing behavior visualization covers moving the kettle, turning the bottom knob, moving the slider, and turning on the oven light.