Policy Improvement using Human Interventions
Results at a Glance
- In our experience, imitation learning for high-dexterity tasks hits a wall, and neither larger datasets nor longer training seem to help. Human interventions at failure points during policy rollouts (HG-DAgger-style), followed by some additional training steps, have been surprisingly effective!
- After two iterations of policy improvement, a pi0.5 policy was able to place a bead on a string 50% of the time, compared to 10-20% of the time without policy improvement.
- After only one iteration of policy improvement, a pi0.5 policy was able to insert a tie wrap tip 25% of the time, and get close another 50% of the time. Without interventions, the tie wrap tip was never inserted successfully.
- After one iteration of policy improvement using human interventions, on a dataset of lids and pans, a pi0.5 policy was able to pick up even small lids 95% of the time. It also placed lids correctly on pans 30% of the time, and got close an additional 65% of the time. This compares to 66% pickup, 15% correct placement, and 50% close placement without policy improvement.
- For every task, only 10-15K additional training steps on the augmented datasets were needed following 40K steps of initial pi0.5 training. So this approach was very training-efficient, with no need to retrain from scratch!
Results Table
| Task | pi0.5 | pi0.5 + PI1 | pi0.5 + PI2 |
|---|---|---|---|
| Tie-wrap inserted | 0% | 25% | — |
| Tie-wrap in+close | 40% | 75% | — |
| Bead on string | 10-20% | 10-50% | 50% |
| Bead on+close | 50% | 90-100% | 100% |
| Place lids on pans | 15% | 30% | — |
| Place on+close | 65% | 95% | — |
| Pickup lids | 66% | 95% | — |
HG-DAgger
- Human intervention: Imitation learning (behavioral cloning) can work surprisingly well for initial policy training, but can result in fragile policies. This is because the distribution of states encountered by real robots includes many states/situations not in the imitation learning dataset D0. In other words, the robot following the imitation learning policy, p0, may find itself on a trajectory different from any it was trained on. This can cause the robot to fail to accomplish its task. One solution is to have a human intervene when the robot, following p0, is about to fail. The human intervenes by stopping the robot mid-trajectory and then continues by teleoperating the robot. Such episodes with human interventions can be recorded and saved as a new dataset D1. This new dataset contains states/situations not in the original imitation learning dataset D0, so it can be used to teach the robot how to behave in a greater variety of situations. To improve the initial policy, p0, one can combine D0 and D1 and resume training p0 to obtain a more robust policy p1. This process can then be iterated.
- HG-DAgger: We mostly follow the method developed in HG-DAgger: Interactive Imitation Learning with Human Experts. The HG-DAgger approach itself is a more practical implementation of the fundamental ideas in the original DAgger paper in which the authors prove the benefit of using the current policy to explore the state space while having an expert replace the current policy's actions in that space with expert actions. In DAgger, all of the rollout actions are replaced in this way which creates a very rich dataset with correct behavior over a wide range of states/situations. Unfortunately, replacing rollout actions with expert actions at all rollout states is not practical for real robots. HG-DAgger, instead, intervenes at a rollout state where failure is about to occur, replacing the entire failure segment with a human teleoperated segment. Also, in original DAgger, the new policy pi+1 is trained from scratch. On the other hand, we train the next policy pi+1 by initializing with the policy pi, and we have found this incremental model improvement is very efficient, requiring relatively few additional training steps per iteration. As in both papers, at every iteration, we train on the combination of all datasets: \(\sum_{j=0}^{i+1} D_j\).
- Number and length of interventions: In HG-DAgger, the robot starts with a policy rollout and runs until a human intervenes and teleoperates the robot through a failure segment. Once past the problem segment, the rollout is resumed, and a human can intervene again as often as necessary. All of the teleoperated segments are recorded and added to the dataset. Primarily for simplicity, we have modified this protocol in two slightly different ways: first, we have been teleoperating after the very first intervention all the way to the end of the episode. This is the approach we used for all of the pi0.5 experiments here. Although we did this initially because it required fewer code mods, there is a principled argument in its favor: the teleoperated state at the end of the failure segment may be outside the state distribution used to train the policy. Hence, teleoperating to the end may teach the robot a necessary follow-up to the correction segment. This may be a different way to deal with the out-of-distribution problem that led to rule 2 in RaC.
Our second modification is again to intervene only once but to terminate the episode after teleoperating through the failure region only. This is closer to RaC rule 2. Unlike HG-DAgger, only the first failure segment is recorded (plus some rollout padding, see below). We have been trying this second approach in our GR00T experiments. We also plan, as a third approach, to try multiple intervention segments, as in HG-DAgger. Which approach is best is an open question for us. What is interesting is that after fixing only the first failure and then post-training the policy, on the next iteration the robot typically gets through this failure segment on its own. It then moves on to the next problem region where the next set of interventions deals with that next problem and so on. This may be less efficient than the original HG-DAgger but it does seem to succeed.
- Rollout padding: In HG-DAgger only the teleoperated segments are added to the new dataset, and no rollout segments are added. It is at least implied that imitating rollout segments would be harmful, since it is the policy after all that went off track. However, HG-DAgger did not discuss action-chunking, and there is a reason to keep at least an action-chunk's length of pre-intervention rollout: After training on the dataset augmented with interventions, the new policy, during rollout, might be called anywhere inside a chunk-length window (actually the window is the horizon, the number of steps before the policy is called again, but we set this equal to the chunk-length) prior to an upcoming 'failure'. Hence it would be good for the model to have been trained on all of the states in that window prior to intervention. The incorrect rollout actions in that window come along for the ride, but in our experiments these do not seem to cause a problem. We do not have enough data to determine whether this small rollout padding is helpful or harmful. In our GR00T experiments we set the window = 1 second (tpad=1.0) which corresponds approximately to a chunk size = 32, at 30 fps.
For the pi0.5 experiments described here, we naively used the entire start-to-intervention rollout. While we don't see direct evidence of harm from using the whole rollout in our pi0.5 experiments, we did see degradation in performance when we at first tried this approach with GR00T, so we believe it is the wrong thing to do and in future we plan to use the chunk-length rollout window for pi0.5 too.
Pi0.5
- Place beads on a string: Fig 2. We applied HG-DAgger to improve the performance of our pi0.5 policy trained on the teleoperation bead-on-a-string dataset D0, which contains 50 episodes. To begin, 50 episodes of human interventions were performed to build dataset D0+D1. The initial policy, which had been trained for 40K steps, was trained for an additional 5K steps on D0+D1. This process was repeated to build D2 and train for another 10K steps on D0+D1+D2. Two types of interventions were performed: one to help the robot pick up the bead, and a second to help the robot place the bead on the string. The improved policy is shown running in Fig 2. In one experiment, the percentage of full task completions went from 20% to 50%. Moreover, the percentage of times the robot either placed the bead correctly or got close improved from 50% to 100%. In addition, there was a large improvement in the percentage of times the robot picked up the bead, going from 50% to 90%.
- Close tie wrap: Fig 3. A pi0.5 policy was trained for 40K steps on our tie-wrap dataset D0, which has 50 episodes. This policy was never successful at inserting the tie wrap tip into the head, although it was able to get close about 40% of the time. One set of 50 episodes with human interventions was added to build D0+D1. The original 40K step policy was trained on D0+D1 for an additional 15K steps. Two types of interventions were performed, one to improve how the right gripper grabs the tie wrap, and a second to fix misalignment of the tip just before insertion. The improved policy was then able to insert the tip into the head 25% of the time, see Fig 3! It was able to insert or get close 75% of the time. More iterations are planned.
- Place lid on pan, before and after: Figs 4,5,6. We applied this approach to improve an imitation learning pi0.5 policy trained on our place-lids dataset D0, which has 50 episodes created using teleoperation. We then added another 50 episodes using human interventions, to create a combined dataset D0+D1. There are two types of error the human corrected. First, as shown in Fig 4b and Fig 5b the robot may fail to pick up the lid. By intervening just before this failure, the lid is picked up and then placed on the pan. Second, the robot picks up the lid correctly but does not place it well on the pan, as shown in Fig 6b, in which case the human intervenes just before the lid is misplaced. Starting from the original pi0.5 policy, which had been trained for 40K steps, training was restarted and continued for another 10K steps on D0+D1. The improved performance is shown in Figs 4a, 5a, and 6a. In one test of multiple lids and pans, the robot went from picking up 66% of the lids to picking up 95%. Placement of lids that stay on their pans went from 15% to 30%. Placement of lids that either stay on their pans or that are close to staying on went from 65% to 95%. Only one iteration was performed, but more are planned.
Implementation on our Robot
Fig 7 shows our implementation of policy rollout with human intervention. This is done in the examples/trossen_ai/record.py function in our openpi fork, which also saves episodes in the required lerobot dataset format. This record function runs the current pi0.5 (or another policy) up until the down arrow key is pressed, at which point the robot arm is frozen. In Fig 7, this happens just before the robot attempts to pick up the lid, which it would fail to do. Next the leader arm is sent to the same position as the frozen follower arm. Once the leader arm is in place, pressing the down arrow key again puts the robot arms into teleoperation mode, and the person completes the episode. Notice, Fig 8, that the recorded dataset video smoothly splices together the rollout and teleoperated trajectories. The record.py script also implements 'early exit', 'rerecord episode', and 'stop recording' as in control_robot.py in lerobot. In our Isaac GR00T N1.7 implementation we also allow for multiple interventions per episode and add a tpad variable to control the length of rollout padding before intervention. These options will be added to our openpi fork too in the future.