TinyFM — Training a 5MB Transformer Policy that runs entirely in your browser

Flow Matching for Robotic Imitation Learning — With DCT-Compressed Trajectories

If you’ve been following robotics and imitation learning research lately, you’ve probably seen diffusion policies everywhere. They work remarkably well — but they’re slow. Inference requires dozens of denoising steps, which is painful for real-time control. Flow matching is a cleaner, faster alternative that’s been gaining traction, showing up in state-of-the-art robot foundation models like π0 from Physical Intelligence.

In this post, we’ll build a flow-matching transformer policy for the Push-T benchmark. Along the way we’ll introduce two extensions that make this project more interesting: DCT-compressed trajectories (inspired by π0) to shrink the action representation, and an ONNX export pipeline that gets the model under 5 MB for in-browser inference. My goal is not to make the new record-setting model for this benchmark; instead it is to see if we can get comparable results while scaling down instead of up. The full Colab notebook is linked below, and you can try the live demo in your browser here.

• TinyFM Notebook
• Live Browser Demo

Example simulation run from the live browser demo

If you are unfamiliar with transformers, this article may be of use:

[embed]Transformers on Human Action Recognition (From Scratch!) Creating and evaluating the validity of an encoder-decoder transformer on a multi-modal capable classification task.medium.com

The Environment: Push-T

Push-T is a 2D manipulation benchmark where a circular “agent” (think robot end-effector) must push a T-shaped block into a goal configuration. The state is a 5-dimensional vector: agent position (x, y), block position (x, y), and block rotation angle θ. Actions are 2D target positions for the agent. It’s deceptively simple but gives policies a meaningful planning challenge — you have to navigate around the T to push it correctly, not just poke at it directly.

We train on the pusht_cchi_v7_replay dataset, which contains roughly 200 human demonstrations recorded via teleoperation.

Why Flow Matching?

Diffusion models learn to reverse a noising process — they iteratively denoise random Gaussian noise into a structured output. This works, but because the denoising score function is implicitly defined through a forward diffusion chain, inference requires many steps (typically 50–100) to produce a clean sample.

Flow matching takes a more direct approach. Instead of learning a denoising score, it learns a velocity field that transports samples from a noise distribution to the data distribution along straight interpolated paths.

Given a target trajectory X_1​ and a noise sample x_0 ~ N(0,I), we define the linear interpolant at time t∈[0,1]:

Linear interpolant w.r.t. t

The velocity along this path is constant:

The model learns to predict this velocity at any point along the interpolated path, conditioned on the current observation:

Because the optimal transport paths are straight lines — not the curved, multi-step trajectories that diffusion requires — the learned field can be integrated accurately with far fewer steps. In practice, 10–20 steps gives you clean samples at inference time, compared to 50–100 for diffusion. The training loss is also simpler and more numerically stable than score matching.

Extension: DCT-Compressed Trajectories

Here’s where things get interesting. Standard diffusion and flow-matching policies predict action trajectories timestep-by-timestep: a model with prediction horizon T=16 outputs 16 action vectors, one per step. The model has to operate over a T×2=32-dimensional action space.

π0, Physical Intelligence’s robot foundation model, uses a trick from signal processing: represent trajectories in the frequency domain using the Discrete Cosine Transform (DCT). Instead of predicting all 16 timesteps directly, you predict the kmost significant DCT coefficients of the trajectory, then reconstruct the full trajectory with an inverse DCT at inference time.

Why does this work? Trajectories are smooth. A robot arm’s path from A to B doesn’t have sharp discontinuities — it’s band-limited in frequency space. The first few DCT coefficients capture the dominant structure (overall direction, curvature), while the high-frequency coefficients mostly capture noise anyway. You can throw those away.

from scipy.fftpack import dct, idct

def dct_compress(traj, k):
    """Compress a trajectory to its first k DCT coefficients."""
    return dct(traj, axis=0, norm='ortho')[:k]
def dct_decompress(coeffs, t):
    """Reconstruct a full trajectory from DCT coefficients."""
    pad = np.zeros((t, coeffs.shape[1]))
    pad[:coeffs.shape[0]] = coeffs
    return idct(pad, axis=0, norm='ortho')

In our implementation we use k=10 coefficients to represent 16-step trajectories — compressing the action representation by 37.5%. The model operates entirely in DCT coefficient space: both its training targets and inference outputs are coefficient vectors, and we only invert back to trajectory space at the very end.

This has a concrete benefit for deployment: the flow-matching action head is predicting 10 vectors of dimension 2, not 16. The lower-dimensional output space means the model has less to learn, converges faster, and runs fewer ONNX inference calls per Euler integration step in the browser.

Observation Preprocessing: Handling Angles

One subtle but important detail. The raw Push-T state includes a block rotation angle θ. Feeding raw angles to a neural network is problematic — there’s a discontinuity at ±π where numerically similar angles like 3.14 and −3.14 are far apart in value, even though they represent the same physical orientation.

The fix is standard: replace θ with its sine and cosine.

def angle_to_sin_cos(angle):
    return np.stack([np.sin(angle), np.cos(angle)], axis=-1)

This expands our 5-dimensional state to 6 dimensions (agent x/y, block x/y, sin θ, cos θ), giving the model a smooth and continuous representation of rotation with no discontinuities. The normalized positional features are concatenated with these trig features so we get the best of both: normalized spatial context and wrapping-safe angle representation.

The Model

The full model is an encoder-decoder transformer with a flow-matching action head. The encoder processes the observation history; the decoder refines noisy DCT coefficient tokens into clean ones, conditioned on the encoded observations.

class FlowMatchTransformer(nn.Module):
    def __init__(self, obs_dim, dct_k, d_model, nhead, num_layers, dropout,
                 max_obs_len=16, max_dct_len=16):
        super().__init__()
        self.dct_k = dct_k

        # Observation encoder
        self.obs_embed = nn.Linear(obs_dim, d_model)
        self.obs_pos = nn.Parameter(torch.randn(max_obs_len, d_model) * 0.02)

        # Noisy trajectory token embedding
        self.x_embed = nn.Linear(2, d_model)
        self.x_pos = nn.Parameter(torch.randn(max_dct_len, d_model) * 0.02)

        # Flow timestep embedding (t ∈ [0, 1])
        self.time_mlp = nn.Sequential(
            nn.Linear(1, d_model),
            nn.SiLU(),
            nn.Linear(d_model, d_model)
        )

        encoder_layer = nn.TransformerEncoderLayer(
            d_model=d_model, nhead=nhead, dropout=dropout, batch_first=False)
        self.encoder = nn.TransformerEncoder(encoder_layer, num_layers=num_layers)

        decoder_layer = nn.TransformerDecoderLayer(
            d_model=d_model, nhead=nhead, dropout=dropout, batch_first=False)
        self.decoder = nn.TransformerDecoder(decoder_layer, num_layers=num_layers)

        self.out = nn.Linear(d_model, 2)

    def forward(self, obs, x_t, t):
        # Encode observation history into memory
        obs_emb = self.obs_embed(obs) + self.obs_pos[:obs.shape[1]].unsqueeze(0)
        memory = self.encoder(obs_emb.permute(1, 0, 2))

        # Embed noisy DCT tokens, condition on flow timestep
        time_emb = self.time_mlp(t).unsqueeze(1)
        x_emb = self.x_embed(x_t) + self.x_pos[:x_t.shape[1]].unsqueeze(0) + time_emb
        dec = self.decoder(x_emb.permute(1, 0, 2), memory)

        return self.out(dec.permute(1, 0, 2))

A few design choices worth highlighting:

Encoder processes the last 2 observation timesteps into a contextual memory via multi-head self-attention. The structure mirrors the encoder from the HAR post, but rather than pooling for classification, it passes memory tokens to the decoder via cross-attention.

Time MLP with SiLU: The flow timestep tis projected through a small 2-layer MLP and added to every decoder token embedding. SiLU (Sigmoid Linear Unit) works better here than ReLU for smooth timestep conditioning — it avoids dead neurons on continuous-valued inputs and produces smoother gradients across the [0,1]range.

Decoder tokens are DCT coefficients: Each decoder token here is a 2D vector (one DCT coefficient per spatial axis). The decoder refines all k=10 of them in parallel, conditioned on the encoder's observation memory via cross-attention.

Model size: With d_model=128, 2 encoder layers, 2 decoder layers, and 4 attention heads, the total parameter count is around 1.3M — under our 5 MB target in fp16.

Flow Matching Loss and Sampling

Training and inference are both clean to implement, which is one of the main practical appeals of flow matching over diffusion.

Training: sample a random t∈[0,1], linearly interpolate between noise and target, and supervise the predicted velocity.

def flow_loss(model, obs, x1):
    bsz = obs.shape[0]
    x0 = torch.randn_like(x1)
    t = torch.rand(bsz, 1, device=x1.device)
    x_t = (1 - t.unsqueeze(-1)) * x0 + t.unsqueeze(-1) * x1
    v_target = x1 - x0
    v_pred = model(obs, x_t, t)
    return ((v_pred - v_target) ** 2).mean()

Inference: start from pure Gaussian noise, then integrate the velocity field forward with Euler steps.

@torch.no_grad()
def sample_flow(model, obs, steps=100):
    x = torch.randn(obs.shape[0], model.dct_k, 2, device=obs.device)
    dt = 1.0 / steps
    for i in range(steps):
        t = torch.full((obs.shape[0], 1), i * dt, device=obs.device)
        v = model(obs, x, t)
        x = x + v * dt
    return x

This is a forward Euler ODE integration — the simplest possible solver. You can use higher-order methods (Runge-Kutta, Heun’s method) for better accuracy with fewer steps, but Euler at 100 steps works reliably for this task. The output x is a batch of DCT coefficient tensors, which we decompress with dct_decompress and unnormalize to get the final action trajectory in simulator coordinates.

Training Configuration

pred_horizon   = 16   # total trajectory length predicted
obs_horizon    = 2    # past observations fed to encoder
action_horizon = 8    # steps executed before re-planning
dct_k          = 10   # DCT coefficients kept
d_model        = 128
nhead          = 4
num_layers     = 2
dropout        = 0.1
batch_size     = 64
num_epochs     = 200
lr             = 1e-3
flow_steps     = 100  # Euler steps at inference

The obs_horizon / pred_horizon / action_horizon setup follows the receding-horizon paradigm from the diffusion policy paper: at each control step we feed the last 2 observations, predict 16 steps ahead, execute only the first 8 of them, then re-plan. This balances reactivity against planning stability.

Training runs for 200 epochs on an 80/20 train/val split with AdamW.

Exporting for the Browser

Getting a PyTorch model into the browser in 2026 is more tractable than it used to be. The pipeline is:

1. Export the trained model to ONNX
2. Ship the .onnx file as a static asset
3. Run inference in the browser with ONNX Runtime Web (WebAssembly for CPU, WebGPU for GPU)

We export two variants:

# fp32 — for WebAssembly / CPU inference (~9.9 MB)
torch.onnx.export(
    model.float().cpu(),
    (dummy_obs, dummy_x, dummy_t),
    'tiny_flowmatch.onnx',
    input_names=['obs', 'x', 't'],
    output_names=['v'],
    opset_version=18,
)

# fp16 - for WebGPU (~5 MB)
torch.onnx.export(
    model.half().cpu(),
    (dummy_obs_fp16, dummy_x_fp16, dummy_t_fp16),
    'tiny_flowmatch_fp16.onnx',
    input_names=['obs', 'x', 't'],
    output_names=['v'],
    opset_version=18,
)

Both fit comfortably on a CDN-hosted static site with no server-side inference required.

The browser demo runs the full flow-matching loop in JavaScript: initialize from Gaussian noise, call the ONNX model iteratively (reduced to 20 Euler steps in the browser for speed), decompress with a JavaScript DCT, and render the predicted action trajectory overlaid on the Push-T canvas in real time. Normalization statistics (stats.json) and model config are serialized separately and loaded alongside the weights.

Evaluation

Rollout Performance

We evaluate by running closed-loop rollouts in the Push-T simulator for up to 300 steps per episode, measuring mean coverage reward — the fraction of the goal T-shape’s area covered by the pushed block.

The model performs well on favorable initial configurations: states where the block starts near the goal, or where a single pushing motion suffices to achieve significant coverage. In these cases, rewards above 0.8 are common and the rollouts look qualitatively smooth. The agent approaches, makes contact with the correct face of the T, and pushes it into place.

Degenerate Cases

Here’s where it gets honest. The model exhibits a consistent failure mode: oscillating in place when the block is positioned such that pushing it directly would move it further from the goal. In these cases, the agent will approach, push slightly, back off, re-approach from roughly the same angle, and repeat indefinitely. It never discovers the globally correct strategy — navigating around the T to approach from the other side.

This failure mode is not unique to our model. The reference diffusion policy implementation and the LeRobot pretrained flow-matching policy both exhibit similar behavior on Push-T. Looking at the failure configurations across all three, a pattern emerges: they’re all initial states that are underrepresented in the demonstration data. Human demonstrators were generally competent and rarely ended up needing to execute complex block re-engagement maneuvers, so those states simply don’t appear in the training distribution. The policy has no ground truth to imitate there, and behavior cloning gives it nothing to fall back on.

This is the fundamental ceiling of pure imitation learning: your policy is only as good as your demonstration coverage. The larger LeRobot model achieves higher average coverage rewards — but it doesn’t eliminate the degenerate cases, it just encounters them less frequently due to better generalization from more capacity and data augmentation. The failure mode similarity across model scales suggests this is a data problem, not a model capacity problem.

One principled fix is DAgger (Dataset Aggregation): run rollouts with the learned policy, identify states where it fails, collect corrective human demonstrations in those states, and retrain. This iteratively expands coverage into the hard regions of state space. Online RL fine-tuning is another avenue. Both are natural next steps.

Conclusion

We’ve built a complete flow-matching transformer policy for Push-T, with two practical extensions beyond the baseline:

DCT trajectory compression reduces the action representation from 32 dimensions (16 steps × 2) down to 20 (10 DCT coefficients × 2), exploiting the smoothness of robot trajectories to shrink the space the model has to reason over. This is the same compression idea used in π0, and it translates directly into a smaller model that exports cleanly.

ONNX export gets the full model — weights, normalization stats, and config — under 5 MB for deployment as a static web asset. No server, no API, no GPU required on the client side. The fp16 WebGPU path brings latency down further on capable hardware, making real-time interaction in the browser feasible.

The remaining challenge, as with diffusion policies and most imitation learning approaches, is data coverage. Flow matching gives us a cleaner and faster training objective than diffusion — simpler loss, fewer inference steps, more stable gradients. But the hard constraint remains: the policy can’t generalize to states the demonstrations never covered. Closing that gap is where the most interesting work is happening next, whether through DAgger, online RL fine-tuning, or synthetic data generation.

The full training notebook is available here, and you can play with the live browser demo here.

Thanks for reading!