Trajectory Queue

Bases: Generic[Experience]

Pure functions defining the trajectory queue. This queue assumes batches added to the queue are a pytree with a shape prefix of (batch_size, trajectory_length). Consecutive batches are then concatenated along the second axis (i.e. the time axis). During sampling this allows for trajectories to be sampled - by slicing consecutive sequences along the time axis.

Attributes:

Name	Type	Description
init	Callable[[Experience], TrajectoryQueueState[Experience]]	A pure function which may be used to initialise the queue state using a single timestep (e.g. (s,a,r)).
add	Callable[[TrajectoryQueueState[Experience], Experience], TrajectoryQueueState[Experience]]	A pure function for adding a new batch of experience to the queue state.
sample	Callable[[TrajectoryQueueState[Experience]], Tuple[TrajectoryQueueState, TrajectoryQueueSample[Experience]]]	A pure function for sampling a batch of data from the replay queue, with a leading axis of size (sample_batch_size, sample_sequence_length). Note sample_batch_size is the same as the batch size when adding data to the queue. However, sample_sequence_length may be different to the sequence length of data added to the state using the add function.
can_sample	Callable[[TrajectoryQueueState[Experience]], Array]	Whether the queue can be sampled from.
can_add	Callable[[TrajectoryQueueState[Experience]], Array]	Whether the queue can be added to.

See make_trajectory_queue for how this container is instantiated.