rlightning.buffer.utils

class rlightning.buffer.utils.BufferView(data: MockModule('tensordict.TensorDict'))[source]

Bases: object

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

class rlightning.buffer.utils.EpisodeTable(num_storages: int, env_ids: Sequence[str] | None = None, num_train_workers: int | None = None, component_distribution: Dict[str, Dict[str, Dict[str, Any]]] | None = None, node_affinity_env: bool = False, node_affinity_train: bool = False)[source]

Bases: object

Track episode -> storage shard mapping and storage -> train worker mapping.

The MVP version uses a simple, even distribution strategy for both: - env_ids are assigned to storage shards in round-robin order with load balance. - train workers are assigned to storage shards in contiguous, even blocks.

When node affinity is enabled, envs and train workers are bound to storages on the same node using the component distribution view.

get_env_to_storage() → Dict[str, int][source]

Return env -> storage mapping.

get_envs_for_storage(storage_idx: int) → List[str][source]

List env_ids currently mapped to the given storage shard.

get_storage_idx_for_env(env_id: str) → int[source]

Return storage index for env_id, assigning it if unseen.

get_storage_to_train_workers() → Dict[int, List[int]][source]

Return storage -> train worker mapping.

get_train_workers(storage_idx: int) → List[int][source]

Return train workers assigned to the given storage shard.

register_envs(env_ids: Iterable[str]) → None[source]

Register a batch of env_ids with the current assignment strategy.

class rlightning.buffer.utils.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)[source]

Bases: object

Actual storage for data buffer to support shard or unshard storage

add_data_async(item: EnvRet | PolicyResponse) → None[source]

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[source]

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[source]

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[source]

Clear all stored data.

get_data()[source]

Return all stored data as a TensorDict.

get_size() → int[source]

Get the size of the storage

print_timing_summary(reset: bool = False) → None[source]

Print the timing summary of the storage.

property size: int

Get the size of the storage

truncate_episodes(env_ids: List[str]) → None[source]

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[source]

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]

Submodules