rlightning.buffer.utils.storage

Storage backend for buffer data and episode handling.

class rlightning.buffer.utils.storage.ActiveEpisodeBuffer(env_meta_list: List[EnvMeta], auto_truncate_episode: bool, postprocess_fn: Callable, device: MockModule('torch.device'))

Bases: object

A helper class to manage active episode buffers for multiple environments.

add_transition(env_id: str, transition: Dict[str, Any]) → None

Add a transition to the episode buffer for a given environment ID.

pop(env_id: str) → List[MockModule('tensordict.TensorDict')]

Pop and return a list of episodes correlated to an env tagged with given environment ID.

Parameters:

env_id (str) – Environment id.

Returns:

A list of episode

Return type:

List[TensorDict]

pop_done_episode(env_id: str | None = None) → Dict[str, List[Any]] | List[Dict[str, List[Any]]]

Pop and return the done episode buffer for a given environment ID.

class rlightning.buffer.utils.storage.BufferView(data: MockModule('tensordict.TensorDict'))

Bases: object

A read-only, buffer-like view of a TensorDict for dataset creation.

class rlightning.buffer.utils.storage.DataContainer(capacity: int, mode: Literal['circular', 'fixed'], unit: Literal['transition', 'episode'], device: str | MockModule('torch.device'))

Bases: object

clear() → None

Clear the storage.

push(items: List | MockModule('tensordict.TensorDict') | Dict[str, Any]) → None

Push items into the storage.

class rlightning.buffer.utils.storage.Storage(capacity: int, mode: Literal['circular', 'standard'], unit: Literal['transition', 'episode'], env_meta_list: List[EnvMeta], device: str | MockModule('torch.device'), obs_preprocessor: Callable, reward_preprocessor: Callable, env_ret_preprocess_fn: Callable, policy_resp_preprocess_fn: Callable, preprocess_fn: Callable, postprocess_fn: Callable, auto_truncate_episode: bool)

Bases: object

Actual storage for data buffer to support shard or unshard storage

add_data_async(item: EnvRet | PolicyResponse) → None

Add a single EnvRet or PolicyResponse to the storage. It is usually used in async rollout that cannot collect env_ret and policy_resp pair at the same time.

Parameters:

item (Union[EnvRet, PolicyResponse]) – The EnvRet or PolicyResponse to add

add_episode(episode: Dict | MockModule('tensordict.TensorDict'), num_envs: int = 1) → None

Add a complete episode to the storage.

Parameters:

• episode (Dict | TensorDict) – The episode to add.

• num_envs (int) – Number of parallel environments represented in the episode.

add_transition(env_ret: EnvRet, policy_resp: PolicyResponse) → None

Add a transition with an env_ret and policy_resp pair to the storage.

Parameters:

• env_ret (EnvRet) – The environment return.

• policy_resp (PolicyResponse) – The policy response.

clear() → None

Clear all stored data.

get_data()

Return all stored data as a TensorDict.

get_size() → int

Get the size of the storage

print_timing_summary(reset: bool = False) → None

Print the timing summary of the storage.

property size: int

Get the size of the storage

truncate_episodes(env_ids: List[str]) → None

manually finish current episode to actually save into replay buffer’s storage. :param env_ids: list of env ids to truncate :type env_ids: List[str]

truncate_one_episode(item: str | Any) → None

manually finish current episode to actually save into replay buffer’s storage. :param item: env_id or object with env_id attribute :type item: Union[str, Any]