unienv_data.batches.slicestack_batch¶
SliceStackedBatch(batch: BatchBase[BatchT, BArrayType, BDeviceType, BDtypeType, BRNGType], fixed_offset: BArrayType, get_valid_mask_function: Optional[Callable[[SliceStackedBatch, BArrayType, BatchT, Dict[str, Any]], BArrayType]] = None, fill_invalid_data: bool = True, stack_metadata: bool = False)
Bases: BatchBase[BatchT, BArrayType, BDeviceType, BDtypeType, BRNGType]
A batch that stacks frames with given fixed offsets. This is a read-only batch, since it is a view of the original batch. If you want to change the data, you should mutate the containing batch instead.
expand_index(index: BArrayType) -> BArrayType
Expand indexes to slice the data Args: index (BArrayType): A 1D tensor of indices to expand. Can be boolean or integer. Returns: BArrayType: A 2D tensor of indices, where each row corresponds to the expanded indices for each batch item.
get_valid_mask_flattened(expanded_idx: BArrayType, data: BArrayType, metadata: Dict[str, Any]) -> BArrayType
Get the valid mask for the flattened data. Args: expanded_idx (B, T): A 2D tensor of indices to slice the data. data (B, T, *D): The data to slice. Returns: Valid Mask (B, T)
get_valid_mask(expanded_idx: BArrayType, data: BatchT, metadata: Dict[str, Any]) -> BArrayType
Get the valid mask for the data. Args: expanded_idx (B, T): A 2D tensor of indices to slice the data. data (BatchT): The data to slice. Returns: Valid Mask (B, T)
fill_data_with_stack_mask(space: Optional[Space[Any, BDeviceType, BDtypeType, BRNGType]], data: Union[BArrayType, BatchT, Any], valid_mask: BArrayType) -> Union[BArrayType, BatchT, Any]
Fill the data with the mask as if the frames were frame-stacked.
data is updated in place where the backend allows it (backend.at(...).set;
JAX and read-only arrays are copied). This is safe because data is the result
of fetching an index array from the wrapped batch, a fresh gather: batches wrapped
by a SliceStackedBatch must not return arrays aliasing their storage from
get_at* / get_flattened_at* with index arrays.
Args:
space (Optional[Space]): The space to fill the data with. If None, the data is assumed to be a backend array.
data (Union[BArrayType, BatchT, Any]): The data to fill sized (B, T, *D)
valid_mask (BArrayType): The mask to fill the data with, sized (B, T)
Returns:
Union[BArrayType, BatchT, Any]: The filled data, sized (B, T, *D)
get_at_with_metadata(idx: Union[IndexableType, BArrayType]) -> Tuple[BatchT, Dict[str, Any]]
get_column(nested_keys: Sequence[Any]) -> BatchBase
Lazy view over one field of the stacked samples.
Stacks self.batch.get_column(nested_keys) instead of stacking whole samples and
dropping the other fields afterwards, so only that column (plus the episode column the
built-in valid-mask helpers read) is gathered. Values, spaces, metadata and the
flattened layout are those of SubItemBatch(self, nested_keys), which is returned
when that is not possible (a custom valid mask function, which may read any field, a
mask read from metadata, or a column view with different metadata).
close() -> None
Release this view. The wrapped batch is not owned by the view and stays open;
close it through whoever created it (BatchBase.__del__ calls close(), so
closing the source here would tear it down whenever a view is garbage-collected).
staticmethod
¶
get_valid_mask_function_with_episodeid_key(episode_id_key: Union[str, int] = 'episode_id', is_in_metadata: bool = False) -> Callable[[SliceStackedBatch, BArrayType, BatchT], BArrayType]
staticmethod
¶
get_valid_mask_function_with_episode_end_key(episode_end_key: Union[str, int] = 'episode_end', is_in_metadata: bool = False) -> Callable[[SliceStackedBatch, BArrayType, BatchT], BArrayType]
set_flattened_at(idx: Union[IndexableType, BArrayType], value: BArrayType) -> None
Overwrite existing samples using flattened data.
append_flattened(value: BArrayType) -> None
Append one flattened sample to the batch.
extend_flattened(value: BArrayType) -> None
Append a batched block of flattened samples.
set_at(idx: Union[IndexableType, BArrayType], value: BatchT) -> None
Overwrite existing samples using structured data.
remove_at(idx: Union[IndexableType, BArrayType]) -> None
Remove one or more samples from the batch.
extend_from(other: BatchBase[BatchT, BArrayType, BDeviceType, BDtypeType, BRNGType], chunk_size: int = 8, tqdm: bool = False) -> None
Copy data from another batch in bounded-size chunks.
get_slice(idx: Union[IndexableType, BArrayType]) -> BatchBase[BatchT, BArrayType, BDeviceType, BDtypeType, BRNGType]
Create a lazy view over a subset of indices.