rlightning.buffer

class rlightning.buffer.BufferConfig(**data: Any)

Bases: Config

Configuration for the data buffer.

auto_truncate_episode: bool = False

Whether to automatically truncate episodes when they are done

capacity: int

Maximum number of transitions the buffer can store

node_affinity_env: bool = False

Whether to enable node affinity for environment workers

node_affinity_train: bool = False

Whether to enable node affinity for training workers

sampler: SamplerConfig | None = None

The sampler configuration

setup_default_sampler = MockModule('pydantic.model_validator')

storage: StorageConfig = MockModule('pydantic.Field')

The storage backend configuration

type: Literal['ReplayBuffer', 'RolloutBuffer']

The type of the buffer

class rlightning.buffer.DataBuffer(config: BufferConfig, obs_preprocessor: Preprocessor | None = <function default_obs_preprocessor>, reward_preprocessor: Preprocessor | None = <function default_reward_preprocessor>, env_ret_preprocess_fn: Callable | None = <function default_env_ret_preprocess_fn>, policy_resp_preprocess_fn: Callable | None = <function default_policy_resp_preprocess_fn>, preprocess_fn: Callable | None = <function default_preprocess_fn>, postprocess_fn: Callable | None = <function default_postprocess_fn>)

Bases: ABC

Data buffer for storing and sampling reinforcement learning experience.

This class provides a unified interface for episode and transition storage, supporting both unified and sharded storage backends, various sampling strategies, and flexible preprocessing pipelines.

add_batched_data_async(batched_data: BatchedData, truncations: List[bool] | None = None) → None

Add batched data which is either batched env ret or batched policy resp to the buffer. This method is useful when async rollout between env and eval policy.

Parameters:

• batched_data (BatchedData) – A batch of environment returns or policy responses.

• truncations (Optional[List[bool]]) – A list indicating whether each episode that corresponds to the given data should be truncated. Defaults to None.

add_batched_transition(batched_env_ret: BatchedData, batched_policy_resp: BatchedData, truncations: List[bool] | None = None) → DataBuffer

Add transition with pairs of env_ret and policy_resp to the buffer. This method will automatically handle the preprocess of given env_ret and policy_resp, add and organize them in an internal episode_buffer, and finally store the episode_buffer to the storage. It is useful when training with in distributed.

Parameters:

• batched_env_ret (BatchedData) – A batch of environment returns.

• batched_policy_resp (BatchedData) – A batch of policy responses.

• truncations (Optional[List[bool]]) – A list indicating whether each episode that corresponds to the given env_ret and policy_resp pair should be truncated. Defaults to False.

• is_eval (Optional[bool]) – Whether the transitions are from evaluation. Defaults to False.

add_data_async(env_id: str, data: EnvRet | PolicyResponse | MockModule('ray.ObjectRef'), truncated: bool | None = False) → None

Add data which is either env ret or policy resp to the buffer. This method is useful when async rollout between env and eval policy.

Parameters:

• env_id (str) – The env id.

• data (Union[EnvRet, PolicyResponse, ray.ObjectRef]) – An environment return or a policy response.

• truncated (Optional[bool]) – Whether the episode is truncated after this data. Default to False.

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

Adds a complete episode to the replay buffer.

Parameters:

• episode (Union[Dict, TensorDict, ray.ObjectRef]) – The episode data to be added to the buffer.

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

add_transition(env_id: str, env_ret: EnvRet | MockModule('ray.ObjectRef'), policy_resp: PolicyResponse | MockModule('ray.ObjectRef'), truncated: bool | None = False) → None

Add transition with a pair of env_ret and policy_resp to the buffer. This method will automatically handle the preprocess of given env_ret and policy_resp, add and organize them in an internal episode_buffer, and finally store the episode_buffer to the storage. It is useful when training with in distributed.

Parameters:

• env_id (str) – The env id.

• env_ret (Union[EnvRet, ray.ObjectRef]) – An environment return.

• policy_resp (Union[PolicyResponse, ray.ObjectRef]) – A policy response.

• truncated (Optional[bool]) – Whether the episode is truncated after this transition. Defaults to False.

• is_eval (Optional[bool]) – Whether the transition is from evaluation. Defaults to False.

clear() → None

Clear all data from the buffer.

get(item: int | Sequence[int] | Dict[str, Any])

Get data from the buffer by sample info.

Parameters:

item – Index, sequence of indices, or dict with storage_idx and indices.

Returns:

TensorDict containing the requested data.

Raises:

TypeError – If item is not a valid type.

get_all() → MockModule('tensordict.TensorDict') | Dict[int, MockModule('tensordict.TensorDict')]

Get all data in the storage.

Warning: This is a debugging function and should be used with caution as it may return large amounts of data.

Returns:

All stored data, or dict mapping storage index to data if sharded.

init(env_meta_list: List[Dict] | None = None, env_ids: List[str] | None = None) → None

Initialize the buffer with environment metadata.

Parameters:

• env_meta_list – List of environment metadata dictionaries.

• env_ids – List of environment identifiers.

• train_worker_num – Number of training workers for data distribution.

init_storage_table(env_ids: List[str] | None = None, train_worker_num=1) → None

Initialize storage table for env -> storage and storage -> train workers mapping.

print_timing_summary(reset: bool = False) → None

Print timing summary for profiling.

Parameters:

reset – If True, reset timing statistics after printing.

sample(batch_size: int | None = None, shuffle: bool | None = True, drop_last: bool | None = True) → List[Dict]

Sample a batch of data from the buffer.

Sharded storage support: - Each storage samples indices based on its own size. - DistributedSampler then splits those indices to the workers bound to that storage.

Returns:

A dictionary mapping worker rank

to a dictionary containing the storage index and indices for the sampled data.

Return type:

worker_sample_info (Dict[int, Dict[str, np.ndarray]])

size() → int

Get the number of transitions stored in the buffer.

Returns:

The number of transitions in the buffer.

Return type:

int

truncate_episodes(items: List[str | Any]) → None

Manually truncate multiple episodes, finalize their collection, and trigger postprocessing and saving to the replay buffer.

Parameters:

items (List[Union[str, Any]]) – List of env_ids or objects with env_id attribute

truncate_one_episode(item: str | Any) → None

Manually truncate a single episode, finalize its collection, and trigger postprocessing and saving to the replay buffer.

Parameters:

item (Union[str, Any]) – env_id or object with env_id attribute

Subpackages

• rlightning.buffer.utils
  • rlightning.buffer.utils.postprocess
  • rlightning.buffer.utils.preprocess
  • rlightning.buffer.utils.preprocessors
  • rlightning.buffer.utils.storage
  • rlightning.buffer.utils.table
  • rlightning.buffer.utils.utils

Submodules

• rlightning.buffer.base_buffer
• rlightning.buffer.replay_buffer
• rlightning.buffer.rollout_buffer
• rlightning.buffer.sampler