Lately some robot companies have claimed that their robots can run at ultra human speed. Such really makes me wonder: how with imitation learning?

Did they hire speedsters to record the tasks? Can you get The Flash on demand?

That seems unlikely and I don't really know how they did that since they obviously didn't share (maybe optimizing splines?). So in this article, I'll try a few fun and naive ways to speed up action chunks.

Sorry in advance, I won't measure the average time rigorously, rather I'll show you the video for you to decide qualitatively. This is an evening project, I just want to experiment with this after dinner and before I go to sleep, which is roughly 6 hours.

My setup

Before the details, a quick note on what I'm running.

My rig is a ReBot arm in the middle of an aluminum extrusion frame. The task that will be used in this article is mentioned in the previous article.

The action chunk policy I use is Diffusion Policy. The reason behind this is simply because I experimented with it extensively for this task in the last article. In addition, I also optimized each pass of the policy itself to run sub-15ms, which allows me to eliminate wait time between action chunks.

Here is the baseline of the task we are "speeding up" today.

The reason I put "speeding up" in quotes is because execution time is a distribution, we can't compare distributions with one sample. In this article, I'll show you the videos for you to decide, but take execution time with a grain of salt (more like a whole bag of salt). Please read this line again in case you want to comment something like "hey bro i think the policies' durations are the same" after reading this article.

Going from first principles

This goes without saying: if you want to do something fast, you just have to do it faster. If you want your policy to be fast, you need it to execute its action chunk faster.

An action chunk is characterized by two things: how many actions it contains, and how long it takes to execute. A horizon-31 chunk has 31 actions. At a 30 Hz control rate, that chunk takes just over a second to play out.

So there are two levers. You can shorten the time the chunk takes by reducing chunk size or increase the control rate, or you can make the actions inside the chunk cover more distance. The first is easy to grasp. The second is trickier, because execution time stays the same and the motion itself has to change.

Let's start with the first.

Subsampling the action chunk

The most obvious way to go faster is to throw away actions. If your chunk has 31 actions, take every other one and play the remaining 16 at the same control rate. This shortens your execution time by half.

Here is that in practice.

When we skip 2 actions for every action, effectively running 3x, we start to see the downside of this method. You can throw away useful actions, which makes the policy unable to finish the task. In addition, the policy works jaggedly, because we throw away actions without regards for their deltas, meaning big actions get discarded the same way as small actions.

Subsampling with a spline

One possible fix here for the jaggedness is to stop treating the chunk as a list of points and start treating it as a curve. You can fit a spline through the predicted actions, in joint space, then resample the spline at whatever density you want.

A spline doesn't necessarily need to be a line in 2d or 3d space, it can be of any dimension you want, which is why you can just do a spline in the joint space of your robot. It does help with the trajectory for the spline to be in end-effector space though.

Here is 2x speed-up using a spline.

You can also make the action chunk slower this way, much like interpolation. Here is 0.5x.

Increasing the control rate

The second obvious method is increasing the control rate of your robot. You can keep all 31 actions in your action chunk and push them out at 60 Hz instead of 30. The chunk finishes in half the time and you don't throw away a single waypoint, so the trajectory is preserved exactly.

The drawback of this method is your robot itself. Can it work at higher control rate? Also, since the policy is trained at 30 Hz, the action deltas are also supposedly in the 30 Hz distribution of deltas. Can your robot follow such deltas at higher control rate? Mine can, at least up to 90 Hz, as I have tested.

Here is doubling the control rate.

Here is tripling the control rate.

Dynamic scaling

These two naive methods have one common downside, which is when you speed up your robot by sending positions only, it’s going to be clumsy. Since you don’t care about force control, and just output positions faster, the robot is going to solve it by giving more torque, which is bad.

Imagine you are trying to nudge a glass bottle by a teeny bit, if the policy is slow, it’s good, but if it’s fast, it’ll break the glass bottle.

So what if the speedup isn't uniform? What if the robot runs flat out across free space and goes slow right before it touched the glass?

You can absolutely do that. In my case, I vary the speed of the robot based on the variations of the gripper. If the upcoming gripper delta is high, meaning it's manipulating an object, go slow. If the upcoming gripper delta is low, go fast.

If you vary the control rate, a pseudo-code of the control loop will look like this:

# pseudocode: variable control rate
while running:
    action = chunk[i]
    send(action)
    scale = speed_scale(obs)        # 1.0 = full speed, 0.2 = crawl
    sleep(base_dt / scale)
    i += 1

Here is it in practice.

If you vary the subsampling, a pseudo-code of the control loop will look like this:

# pseudocode: variable frame skip
t = 0.0
while t < horizon:
    send(spline(t))
    t += stride * speed_scale(obs)  # larger stride = faster

Here is it in practice:

Speeding up your dataset

Everything above fights the problem at inference time. But what if you never touch the control loop at all, and instead train the policy to be fast by speeding up the dataset?

Mechanically this is similar to subsampling the chunk, since you're decimating a trajectory either way, but the implications are different. Everything stays exactly the same, except the actions themselves just cover more ground. The policy learns fast motion as the thing it does, rather than having fast motion imposed on it afterwards, like you learning how to watch a podcast at 2x speed lol.

This also means you have the same replan cost for your inference, because the execution time of a chunk stays the same.
Here is a policy trained on a dataset that is sped up by 2x.

And the results?

From my qualitative experiment, speeding up by increasing control rate yields the highest quality roll-outs. This is pure assumption from watching the policy, but it does look like that. Dynamic scaling on top of that is most impressive because it can decide when to speed up, this in turn make it less prone to errors when handling objects.

I also tested training a policy on different variants of sped-up datasets (1x, 2x, 3x, and 4x) by conditioning it on the speed of the dataset. This means I can effectively tune the speed at run time. The results were horrendous, it couldn't do the task.

Either way, I don't aim to make this a reference on the impact of sampling on task success rates. If I did, I would have been much much more rigorous. Again, take this with a big grain of salt.

A fun note

Actions after all are just a time series, we can treat it like a time series, perhaps like an audio time series. This notion has inspired works like FAST Token or Neural Action Codec.

I think it's fun to treat actions from a robot like an audio recording, in that when sped up or down, it has a pitch. Someone should play the song of a roll-out lol.

By the way, kid you not, changing audio pitch is quite a non-trivial problem that you'll likely solve wrong when first tackling it. Here is a fun video to watch.