Why we need Flow Matching?

Unlike MSE-based policies, which generate an entire action chunk directly and may have difficulty representing complex or multimodal action distributions, flow matching learns a conditional vector field that gradually transforms random noise into a realistic action sequence. The method is closely related to diffusion models, but it is typically simpler to implement and often achieves stronger performance.

Action Chunking with Flow Matching Loss

Let $A_t^{(j)}$ denote the $j$-th ground-truth action chunk, and let $A_{t,0}^{(j)} \sim \mathcal{N}(0,I)$ be Gaussian noise with the same dimensions. During training, we sample a flow-matching time value $\tau^{(j)} \sim \mathcal{U}(0,1)$ and construct an interpolated action chunk:

$$A_{t,\tau}^{(j)} = \tau^{(j)}A_t^{(j)} + \left(1-\tau^{(j)}\right)A_{t,0}^{(j)}.$$

A neural network $v_\theta$ is then trained to estimate the velocity that moves the interpolated sample $A_{t,\tau}^{(j)}$ toward the target action chunk $A_t^{(j)}$. The corresponding flow-matching objective is

$$\mathcal{L}_{\mathrm{FM}}(\theta) = \frac{1}{B} \sum_{j=1}^{B} \left\| v_\theta \left( o_t^{(j)}, A_{t,\tau}^{(j)}, \tau^{(j)} \right) - \left( A_t^{(j)}-A_{t,0}^{(j)} \right) \right\|_2^2.$$

At inference time, the process begins with a random sample $A_{t,0} \sim \mathcal{N}(0,I)$. The action chunk is generated by solving the ordinary differential equation

$$\frac{dA_{t,\tau}}{d\tau} = v_\theta(o_t,A_{t,\tau},\tau)$$

from $\tau=0$ to $\tau=1$. A straightforward way to solve this equation is Euler integration:

$$A_{t,\tau+\frac{1}{n}} = A_{t,\tau} + \frac{1}{n} v_\theta(o_t,A_{t,\tau},\tau).$$

This update is applied $n$ times, where $n$ is the number of integration, or denoising, steps. The resulting sample $A_{t,1}$ is treated as the final action chunk $A_t$, which is then executed open-loop.

Implementation

models.py

class FlowMatchingPolicy(BasePolicy):
    """Predicts action chunks with a flow matching loss."""

    def __init__(
        self,
        state_dim: int,
        action_dim: int,
        chunk_size: int,
        hidden_dims: tuple[int, ...] = (128, 128),
    ) -> None:
        super().__init__(state_dim, action_dim, chunk_size)

        layers: list[nn.Module] = []
        current_dim = state_dim
        for hidden_dim in hidden_dims:
            layers.append(nn.Linear(current_dim, hidden_dim))
            current_dim = hidden_dim
        layers.append(nn.Linear(current_dim, chunk_size * action_dim))
        self.model = nn.Sequential(*layers)

    def forward(self, state: torch.Tensor, action_chunk: torch.Tensor, tau: torch.Tensor) -> torch.Tensor:
        """Evaluate v_theta(state, action, tau)"""

        flat_action = action_chunk.flatten(start_dim=-2)
        network_input = torch.cat(
            [state, flat_action, tau]
        )
        flat_velocity = self.model(network_input)
        return flat_velocity.reshape(-1, self.chunk_size, self.action_dim)



    def compute_loss(
        self,
        state: torch.Tensor,
        action_chunk: torch.Tensor,
    ) -> torch.Tensor:
        # A_{t,0} \sim N(0, I)
        initial_noise = torch.randn_like(action_chunk)

        # flow time \tau. Shape (B, 1, 1) allows broadcasting over time and action dimensions
        batch_size = state.shape[0]
        tau = torch.rand(
            batch_size, 1, 1, 
            device=action_chunk.device, dtype=action_chunk.dtype
        )

        # A_{t,\tau} = \tau * A_t + (1-\tau) * A_{t,0}
        interpolated_action = (
            tau * action_chunk + (1.0 - tau) * initial_noise
        )

        # Compute loss: compare v_\theta and A_t - A_{t,0}
        target_velocity = action_chunk - initial_noise
        predicted_velocity = self.model(
            state=state, 
            action_chunk=interpolated_action, 
            tau=tau.reshape(batch_size, 1)
        )
        return nn.functional.mse_loss(predicted_velocity, target_velocity)

    @torch.inference_mode()
    def sample_actions(
        self,
        state: torch.Tensor,
        *,
        num_steps: int = 10,
    ) -> torch.Tensor:

        batch_size = state.shape[0]

        # Start the ODE at A_{t,0} \sim N(0,I).
        action = torch.randn(
            batch_size, self.chunk_size, self.action_dim,
            device=state.device, dtype=state.dtype
        )

        step_size = 1.0 / num_steps

        # Forward Euler integration from tau = 0 to tau = 1
        for step in range(num_steps):
            tau_value = step * step_size
            tau = torch.full(
                (batch_size, 1), tau_value, 
                device=state.device, dtype=state.dtype
            )
            velocity = self.model(
                state=state, 
                action_chunk=action, 
                tau=tau,
            )
            action = action + step_size * velocity

        return action

Contributors: Junyuan He

Previous pageSegment Anything Model (SAM)

Next pageImitation Learning: Action Chunking with MSE loss