Data flow, tensor shapes, and the code map for this repository

12.1 Rollout buffer record, per macro-step

Field	Shape / type	Producer	Consumer
image	uint8 [3, 224, 448]	worker	learner (cache recompute, imagination)
proprio	float32 [8] (normalised)	worker	learner
task_id, seed, group_id	ints	worker	advantage grouping, context lookup
chain	float32 [K+1, 32, 7]	worker	actor log-prob recompute
k_star	int (or [K_sde] ints for the variant)	worker	which step(s) are stochastic
logp_old	float32 [1] (or [K_sde])	worker	PPO ratio
future_frames	uint8 [2, 3, 224, 448] at +4, +8	worker	reconstruction reward
r_task, done	float, bool	worker	returns
r_rec	float	learner	returns
$R$, A_hat	float	learner	actor loss

Per successful episode, additionally: frames uint8 [T, 3, 224, 448] and executed_actions float32 [T, 7] for video SFT.

12.2 Proposed module layout

All new code goes under src/fastwam/rl/ so the existing training and evaluation paths are untouched.

• flow_sde.py — sample_chain(model, state, K, a, generator, stochastic_steps="one_random") returning (A, chain, k_star, logp_old), used by workers; and logprob_chain(model, states, chains, k_star, a) returning the stochastic-step log-probabilities with gradients, used by the learner. Both must build their σ schedule through model.infer_action_scheduler.build_inference_schedule(K, …) so that the σ and Δ values are bit-identical between rollout and learner.
• rollout_worker.py — a fork of experiments/libero/eval_libero_single.py in which _predict_action_chunk calls sample_chain instead of infer_action, the frames at the required offsets are captured, and episodes are serialised. experiments/libero/worker_pool.py can be reused for GPU scheduling and fault tolerance.
• reward.py — reconstruction_reward(model, batch, use_frozen) running the imagination path and the VAE decode, frame matching, pixel MSE, and a running normaliser.
• advantage.py — grouped returns, group normalisation, degenerate-group dropping, and the hard-task sampling weights.
• losses.py — ppo_flow_loss, feature_kl_loss (with forward hooks capturing the per-layer cache of the live and frozen video expert), and a wrapper that calls model.training_loss with the chosen $\lambda_{\mathrm{action}}$.
• learner.py — the iteration loop of Chapter 11, LoRA injection on model.video_expert (via peft), the two optimisers, checkpointing, and weight broadcast.
• configs/rl/libero_wamrl.yaml — a Hydra config extending sim_libero and libero_optional_idm_2cam224_1e-4, with an RL: block holding the Chapter 11 table.