Guided Policy Search for Stabilizing Contact-rich Motion Plans
Before I get pedantic, let’s be honest about something. This project was all about pushing a box.
Granted, it was really about training an AI model to control a single-point manipulator to push a box across a plane, into a step, and then over the step with exactly one 90-degree topple, followed by finally pushing the box to a goal position, all using non-prehensile manipulation. This work was supported by NSF IIS award #2330794, and took up the first year of me pursuing a PhD. The full published research paper can be found here.
I’m the lead author on this paper, and my co-authors are Chandika Silva (Boise State University), Aykut C. Satici (Boise State University), and Hasan A. Poonawala (University of Kentucky, Lexington).
This page will serve as a combination of the research paper in addition to some editorial notes and commentary after the fact. The kind of stuff that doesn’t fit in research papers.
Abstract: Learning policies for contact-rich manipulation is a challenging problem due to the presence of multiple contact modes with different dynamics, which complicates state and action exploration. Contact-rich motion planning uses simplified dynamics to reduce the search space dimension, but the found plans are then difficult to execute under the true object manipulator dynamics. This paper presents an algorithm for learning controllers based on guided policy search, where motion plans based on simplified dynamics define rewards and sampling distributions for policy gradient-based learning. We demonstrate that our guided policy search method improves the ability to learn manipulation controllers, through a task involving pushing a box over a step.
Editorial notes: Training AL/ML models to control a robot’s actuator to perform non-prehensile manipulation tasks (think pushing, balancing, really anything except grabbing) in contact-rich environments, such as a table with lots of stuff on it or pushing a box up a flight of stairs, is hard. Usually, motion paths are found using simplifications but that often results in more challenges when actually trying it in the real world. We proposed using the motion plan that was developed using the simplifications, and then learning AI/ML models in a more accurate physics simulation. The resulting model would be the new controller for a given task. A natural next step could then be creating more generalized models that apply to more than just one task.
The problem, in English
Say you want a robot to push a box up onto a step, using just one finger touching one point on the box. No grasping allowed, since the box is too big to fit in a gripper anyway. This falls into a category called non-prehensile manipulation, and it's deceptively nasty. Every time the finger's contact point changes, or the box tips, or a face of the box hits the step, the physics governing the system changes too. You go from "pushing a box across a table" to "pivoting a box on its corner" to "a box in free flight for a split second" to "a box landing on a step," and each of those regimes has its own dynamics. Roboticists call these different contact modes, and a controller has to be able to handle all of them, plus the transitions between them, without anyone telling it in advance which mode it's currently in.
The standard way to deal with this is to plan a path first, using a simplified version of the physics (quasi-static or quasi-dynamic, meaning you pretend the object is either always in equilibrium or moving at constant velocity), and then hand that plan to a low-level controller to execute. The problem is that the simplified physics used for planning and the actual physics of the real system don't agree with each other, so a controller that just tries to hit the planned waypoints tends to fall apart the moment reality disagrees with the plan. We saw this ourselves: naively feeding the CMGMP plan (the planner from Cheng et al., 2022) directly into a position controller and letting it run led to the manipulator needing to be manually relocated to keep progressing, since the plan assumed the finger could just teleport its contact point whenever convenient.
Our methods
Instead of treating the simplified motion plan as something to be executed directly, we treated it as a rough draft. A hint. We used it to shape a reward function and a sampling distribution for a
policy gradient algorithm training a neural network controller in a much more accurate, fully dynamic simulator (MJX, JAX's implementation of the MuJoCo physics engine, which is differentiable through contact).
Concretely, the loss function used to train the controller was a weighted sum of terms penalizing distance from the goal, distance from the planned box position, distance from the planned manipulator position, velocity magnitude, and loss of contact between the finger and the box. Turn the "distance from the plan" terms off, and you get an unguided policy that's only trying to reach the final goal, with no idea how to get there. Turn them on, and the plan effectively hands the policy a curriculum: get here first, then here, then here. This is the "guided" in guided policy search.
Training data was generated with a state sampling procedure that specifically targeted the thin sliver of configuration space where the manipulator and box are actually in contact (uniformly sampling the full 2D-rigid-body-times-2D-rigid-body space would waste almost all of its samples on configurations where nothing was touching anything), combined with DAgger, an imitation-learning technique for sampling states along trajectories generated by the current policy rather than only from a fixed distribution. A trimming step discarded any part of a DAgger trajectory once it wandered too far from the region the plan actually covered, which kept bad rollouts from poisoning the training data.
Results
The difference between guided and unguided training wasn't subtle. The unguided policy almost always got stuck in a local minimum, essentially the model equivalent of shrugging, and this
held regardless of the learning rate we tried. The guided policy consistently pushed the box up and over the step to the goal pose. In terms of final position error, the guided controller landed at 0.0305, the unguided controller at 2.5306, and the plan executed directly by a position controller (assuming the finger could freely relocate its contact point) landed at 0.0707, meaning our learned controller actually outperformed the plan it was guided by, without needing to cheat by teleporting the contact point. Training and validation loss curves told the same story: the guided run dropped sharply partway through training as the policy figured out how to get the box onto the step, while the unguided run stayed noisy and flat the whole time.
Stuff that got left out
A few things worth mentioning that don't really belong in a conference paper but are true anyway.
First, "exactly one 90-degree topple" in the intro up top is doing some work. Getting the box to tip cleanly onto the step, rather than sliding, over-rotating, or getting stuck balanced on an edge, was one of those problems where the failure modes were extremely entertaining in simulation and extremely not fun to debug. A lot of the loss function weighting choices in the paper (the relative weights on object position error versus manipulator position error, for instance) came out of iterating on exactly this kind of failure.
Second, this project doubled as a crash course in JAX and differentiable simulation, on top of the control theory and RL side of things. MJX being differentiable through contact was a big part of why we picked it, even though this particular paper doesn't lean on analytic gradients through the simulator directly.
Third, this was very much a proof of concept: one task, one box, one step. The natural next step, which the paper gestures at in its conclusion, is generalizing beyond a single fixed task, and that's an obvious direction for anyone picking this up after us.