factoriax.analysis.recorder

Collect rollout chunks from a training loop into a trajectory.

RolloutRecorder takes the arrays that a collection function already returns and turns them into a Trajectory. A training loop therefore needs one new call, and no change to collect_fn or to any code under jit.

Call record() after each collection call, then finish() once at the end.

The sketch below names collect_fn and total_iters from the training loop it plugs into, so it does not run on its own:

from factoriax.analysis.recorder import RolloutRecorder

recorder = RolloutRecorder(max_episodes=32)
for _ in range(total_iters):
    trajectories, env_states, obs, last_values, _ = collect_fn(...)
    recorder.record(trajectories, env_states)
traj = recorder.finish()

class factoriax.analysis.recorder.RolloutRecorder(max_episodes: int | None = None, record_states: bool = False, state_fields: list[str] = <factory>, _action_chunks: list[ndarray] = <factory>, _reward_chunks: list[ndarray] = <factory>, _done_chunks: list[ndarray] = <factory>, _state_chunks: dict[str, list[~numpy.ndarray]]=<factory>, _num_complete_episodes: int = 0)

Accumulates rollout data across training iterations.

The recorder reads the arrays that a collection function already returns. A training loop therefore needs one new call, and no change to collect_fn or to any code under jit.

Each call to record() stores one (T, N) chunk. finish() joins the chunks and cuts them into episodes at the done flags.

max_episodes

Stop storing chunks once this many episodes have finished, and keep at most this many episodes in the result. None means no limit, and finish() then decides the count.

Type:

int | None

record_states

True to also store the fields named in state_fields from the env states. This costs memory in proportion to the size of those fields, and it is what makes a state-over-time analysis possible.

Type:

bool

state_fields

Which EnvState fields to store when record_states is True.

Type:

list[str]

Notes

The episode count in is_full counts done flags across every parallel environment, and not per environment.

finish(pad_incomplete: bool = True) → Trajectory

Cut the recorded chunks into episodes and build a trajectory.

The chunks are joined along the time axis and then split at the done flags, one environment at a time. State fields are cut at the same boundaries.

Episodes differ in length, so every episode is padded with zeros to the length of the longest one.

Parameters:

pad_incomplete – True to keep the last episode of each environment even when no done flag ended it. False to keep finished episodes only.

Returns:

Actions of shape (B, T_max) or (B, T_max, P), plus rewards, plus any recorded state fields under their trajectory names. B is the episode count after the max_episodes limit.

Return type:

Trajectory

Raises:

ValueError – When no chunk was recorded, or when the split found no episode at all. The second case happens with pad_incomplete=False and no done flag anywhere.

Notes

The padding is zeros, and action 0 is NOOP. Padded steps therefore read as real no-op actions to every analysis function, and the trajectory carries no mask that marks them. A per-episode length is the only way to tell them apart, and this method does not return one.

The method does not clear the buffers. A second call rebuilds the same trajectory. Call reset() to start again.

property is_full: bool

Whether the episode count has reached max_episodes.

Returns:

Always False when max_episodes is None.

Return type:

bool

property num_recorded_steps: int

Total timesteps stored so far, summed over the chunks.

Returns:

The length of the time axis, not multiplied by the number of parallel environments. 0 before the first record().

Return type:

int

record(trajectories: Any, env_states: Any = None) → None

Store one chunk of rollout data.

The call returns at once when the recorder is already full, so a training loop can keep calling it without a guard. The chunk is then dropped.

Parameters:

• trajectories – The struct that collect_fn returns. It must carry action, reward, and done. Actions have shape (T, N) or (T, N, P), and the other two (T, N), where T is the chunk length and N the number of parallel environments. JAX arrays and NumPy arrays both work.

• env_states – The env states from the same call, or None. They are read only when record_states is True. A field named in state_fields but absent from the states is skipped in silence.

reset() → None

Drop every stored chunk and reset the episode count.

The configuration fields keep their values, so the recorder is ready for another run right away.