Understanding Action Chunking with Flow Matching

What is Action Chunking?

Let’s first review the concept of action chunking, popularized by research like 📄Learning Fine-Grained Bimanual Manipulation with Low-Cost Hardware (ACT). In the standard imitation learning, the policy is defined as:

The policy $\pi_{\theta}$ represents a probability distribution over a single action $\mathbf{a}_t$ conditioned on the current observation $\mathbf{o}_t$. Instead of predicting just one single action $\mathbf{a}_t$, the action chunking proposes predicting a sequence, or chunk, of future actions all at once:

where $\mathbf{A}_t=\set{\mathbf{a}_t, \mathbf{a}_{t+1},\mathbf{a}_{t+2},\dots,\mathbf{a}_{t+K-1}}$ is a fixed $K$-length sequence of actions.

As a concrete example, consider the Push-T task from 📄Diffusion Policy: Visuomotor Policy Learning via Action Diffusion, where a robot agent pushes a T-shaped block into a target region.

In this task, the observation relies on low-dimensional state vectors rather than high-dimensional images (i.e., we use a state $\mathbf{s}_t$ instead of a visual observation $\mathbf{o}_t$). The state is a 5-dimensional vector containing:

• $x$-coordinate of the agent
• $y$-coordinate of the agent
• $x$-coordinate of the T-block
• $y$-coordinate of the T-block
• $\alpha$ (rotation angle) of the T-block

The action is a 2-vector $\displaystyle \begin{bmatrix}x\\y\end{bmatrix}$ dictating the agent’s target position.

Why Flow Matching?

The simple answer: generative models like flow matching and diffusion models are more expressive than a mixture of Gaussians.

As an example, in the Push-T environment, a policy powered by a generative model can learn to approach the T-block from the left or from the right.

With a mixture of Gaussians, “the action to the left” averages with “the action to the right” and produces “go straight” — which is not what we want.😅

The Dataset: From Episodes to Training Samples

Expert demonstrations are recorded as episodes of varying lengths, depending on how long it took the expert to solve the task. To avoid jagged arrays during training, all episodes are concatenated into flat arrays. An episode_ends array marks the boundary indices of each episode.

For example, consider two short dummy episodes:

We concatenate their actions and states into flat vectors:

The episode_ends array is [5, 12], meaning Episode 0 spans indices 0-4, and Episode 1 spans indices 5-11. With a chunk_size of $K=3$, a sliding window extracts (state, action_chunk) pairs. Each training sample pairs one state with the next $K$ consecutive actions.

💬Question: Given the episodes above and $K=3$, what are the valid starting indices? Why are indices 3 and 4 excluded for Episode 0?

🗣Answer: The valid starting indices would be [0, 1, 2, 5, 6, 7, 8, 9]. Indices 3 and 4 are excluded because they don’t have $K=3$ actions left before the episode ends. The same logic applies to Episode 1, where index 10 is excluded.

Episode 0:
[a0] [a1] [a2] [a3] [a4]
 ✓    ✓    ✓    ✗    ✗
 0    1    2    3    4

💬Question: During training, the dataloader shuffles and mixes samples from different episodes within a batch. Why is this acceptable?

🗣Answer: Because the underlying task remains constant. Every (state, action_chunk) pair is a self-contained snapshot of an intermediate step toward the same objective. Drawing randomized, uncorrelated transitions across different demonstrations also ensures the batch has high entropy, which improves gradient quality.

Training: What the Network Sees

In Flow Matching, we define a continuous flow from a simple noise distribution to our complex data distribution. Let’s establish our notation. Let the expert action chunk be our target data $\mathbf{A}_1$, and let $\mathbf{A}_0 \sim \mathcal{N}(0, \mathbf{I})$ be a sample of pure noise.

We introduce a flow-time variable $\tau \in [0, 1)$.

During a training step, the network $v_\theta$ receives the current state $\mathbf{s}_t$, a specific flow-time $\tau$, and an interpolated action chunk $\mathbf{A}_\tau$. We construct this straight-line interpolation as:

$\mathbf{A}_\tau = \tau \mathbf{A}_1 + (1 - \tau) \mathbf{A}_0$

The network’s objective is to predict the velocity (the derivative with respect to $\tau$). For a straight-line optimal transport path, the exact target velocity is simply:

For example, if the chunk size $K=3$, then the dimension of the input of the network would be

state(5) + interpolated_chunk(2 * 3) + tau(1) = 12

and the output would be a velocity vector of dimension

action_dim(2) * chunk_size(3) = 6

💬Question: The target velocity ($\mathbf{A}_1 - \mathbf{A}_0$) does not depend on $\tau$. Why does $\tau$ still matter as an input to the network? What would go wrong if we always set $\tau = 0$ during training?

🗣Answer:

The target velocity is constant, but the location $\mathbf{A}_\tau$ in the vector field is constantly changing. We are training the network to predict the correct velocity from any coordinate along the path.

If we only trained on $\tau = 0$, the network would only learn how to predict velocities when looking at pure noise. During inference, after taking the first integration step, the data becomes partially denoised ($\tau > 0$). The network would have no idea how to handle this new, structured input, and the integration would immediately collapse.

💬Question: At $\tau = 0.0$, the network input is pure noise. At $\tau \approx 1.0$, it is almost the exact expert data. At which extreme is the prediction task hardest?

🗣Answer: It is hardest at $\tau = 0.0$. The network is looking at complete static and has to guess the exact trajectory toward a highly specific, structured action chunk.

Inference: Euler Integration

At inference time, we do not have access to expert data. We start by sampling pure noise $\mathbf{A}_0 \sim \mathcal{N}(0, \mathbf{I})$ and use the learned velocity field $v_\theta$ to integrate forward using the Euler method.

If we choose num_steps = 4, our step size is 1 / 4 = 0.25, i.e., $\Delta\tau = 0.25$. The integration process looks like this:

💬 Question: Why is the network never evaluated at $\tau = 1.0$ during inference?

🗣Answer: Because the Euler step taken at $\tau = 0.75$ pushes the chunk exactly to the boundary of $\tau = 1.00$. Once we arrive at the destination, the flow is complete, and we extract the actions to execute.

💬 Question: What happens if you set num_steps = 1? Under what condition would a single step produce a perfect sample?

🗣 Answer: Setting num_steps = 1 means taking one massive Euler step:

This would only produce a perfect sample if the network learned the exact true velocity ($\mathbf{A}_1 - \mathbf{A}_0$) flawlessly. Because Flow Matching paths are straight lines by construction, a perfect velocity prediction means one step is theoretically sufficient. In practice, predictions at $\tau=0$ are the noisiest, so breaking the flow into smaller steps allows the network to correct its course as the chunk becomes more structured.

Reflections

Batch: indices[7], indices[1], indices[5]
     =       8,          1,          6
     → (s8, [a8,a9,a10]),  (s1, [a1,a2,a3]),  (s6, [a6,a7,a8])
       ^^^ from ep1          ^^^ from ep0       ^^^ from ep1