The decision problem (states, actions, rewards, episodes)

RL algorithms are defined over a Markov decision process: a state, an action, a reward, and a transition to the next state. The first design task is to decide what those are for Fast-WAM, and the choices here follow what πRL, SimpleVLA-RL, and RLinf do for chunked VLA policies.

4.1 Time scale: the macro-step

Fast-WAM is queried every 10 environment steps and each query returns 32 actions. The cleanest thing to do is to define one RL time step, which this document calls a macro-step, as one query. The RL "action" is the entire 32×7 normalised chunk the model produced, even though only the first 10 rows are executed. This is the convention used by πRL and SimpleVLA-RL: the policy's probability density is over what the policy actually emitted, and the fact that the environment only consumes part of it is a property of the environment, not the policy.

A LIBERO episode is at most 400 environment steps (700 for LIBERO-10), so at most 40 (or 70) macro-steps.

4.2 State

The state at macro-step $t$ is exactly what infer_action consumes:

• input_image: [1, 3, 224, 448], the two cameras concatenated, built by _obs_to_model_input in eval_libero_single.py.
• proprio: [1, 8], normalised by the dataset statistics.
• context, context_mask: [1, 128, 4096] and [1, 128], the T5 embedding of the prompt. Each LIBERO task has a fixed prompt, so these can be precomputed once per task and shared; the rollout workers then do not need to load T5 at all.

Internally the model turns this into the first-frame VAE latent and then into the per-layer K/V cache. For the purposes of the MDP, the raw image, proprio, and context are the state; the cache is a deterministic function of them.

4.3 Action

The action is the normalised chunk $A \in \mathbb{R}^{32 \times 7}$, the value of latents_action at the end of the denoising loop, before _denormalize_action, the gripper sign flip, and binarisation. Those three operations are deterministic functions of $A$, so the environment's response is a deterministic function of $A$, and the policy density over $A$ is the correct object to differentiate. Do not attempt to define the density over the binarised gripper; you would be trying to differentiate a step function.

4.4 Reward

The reward at macro-step $t$ is the sum of a sparse task term and an optional dense reconstruction term. Chapter 6 defines both. The task term is 1 at the macro-step in which LIBERO's success condition first becomes true and 0 elsewhere.

4.5 Episode termination and discount

An episode ends on success or when the step budget is exhausted. Because the dense reconstruction reward is present at every step, you need a discount factor $\gamma < 1$ to keep returns bounded and to make early good behaviour count more than late; $\gamma = 0.99$ per macro-step is a sensible default, with 0.995 for the longer LIBERO-10 episodes. SimpleVLA-RL and RLinf use effectively $\gamma = 1$ because their reward is a single terminal bit, and if you run the success-only ablation you can do the same.

4.6 Exploration

An imitation-trained deterministic policy never explores. If the sampler always produces the same chunk for the same state and noise, RL has nothing to compare. Exploration comes from two sources: the initial Gaussian noise $x_0$, which is already random, and the per-step noise injected by the stochastic sampler described in the next chapter. The second is the one you control.