rlightning.types¶

Types module for reinforcement learning data structures.

This module provides core data types used throughout the RLightning framework for representing environment returns, policy responses, and batched data.

Available types:

BatchedData: Container for batched environment or policy data.
EnvRet: Environment return data structure.
MultiAgentEnvRet: Multi-agent environment return.
PolicyResponse: Policy action response structure.
MultiAgentPolicyResponse: Multi-agent policy response.
EnvMeta: Environment metadata container.

class rlightning.types.BatchedData(ids: List[Any], data: List[Any | MockModule('ray.ObjectRef')])[source]¶

Bases: object

Container for batched data with associated identifiers.

Holds either actual data or Ray ObjectRefs, maintaining the order of (id, data) pairs. Allows duplicate IDs.

is_future¶: True if any data element is a Ray ObjectRef.

static from_dict(data_dict: Dict) → BatchedData[source]¶

Create BatchedData from a dictionary.

Parameters:: data_dict – Dictionary mapping IDs to data.
Returns:: BatchedData instance.

ids() → Tuple[source]¶

Get all IDs in the batched data.

Returns:: Tuple of all identifiers.

items() → Iterator[Tuple[Any, Any | MockModule('ray.ObjectRef')]][source]¶

Get (id, data) pairs iterator.

Returns:: Iterator of (id, data) pairs.

values() → Tuple[source]¶

Get all data values.

Returns:: Tuple of all data elements.

class rlightning.types.EnvMeta(env_id: str = None, action_space: MockModule('gymnasium.spaces.Space') | None = None, observation_space: MockModule('gymnasium.spaces.Space') | None = None, num_envs: int | None = None)[source]¶

Bases: object

Environment metadata container.

Stores metadata about an environment including its spaces and configuration.

env_id¶

Unique identifier for the environment.

Type:: str

action_space¶

Gymnasium action space.

Type:: MockModule(‘gymnasium.spaces.Space’) | None

observation_space¶

Gymnasium observation space.

Type:: MockModule(‘gymnasium.spaces.Space’) | None

num_envs¶

Number of parallel environments.

Type:: int | None

action_space: MockModule('gymnasium.spaces.Space') | None = None¶: Action space of the environment.

env_id: str = None¶: The environment ID.

num_envs: int | None = None¶: The vectorized number of the environment.

observation_space: MockModule('gymnasium.spaces.Space') | None = None¶: Observation space of the environment.

class rlightning.types.EnvRet(env_id: str, observation: Any, last_reward: float = 0.0, last_terminated: bool = False, last_truncated: bool = False, info: Dict[str, ~typing.Any]=<factory>, _extra: Dict[str, ~typing.Any]=<factory>, ts_env_sent_ns: int = <factory>)[source]¶

Bases: object

Environment return data structure for single-agent interactions.

Represents the return value from environment step() or reset() calls, containing observation, reward, termination status, and additional info.

env_id¶

Unique identifier for the environment instance.

Type:: str

observation¶

Observation after environment step/reset.

Type:: Any

last_reward¶

Reward received after the last step.

Type:: float

last_terminated¶

Whether the episode terminated after last step.

Type:: bool

last_truncated¶

Whether the episode was truncated after last step.

Type:: bool

last_info¶: Additional info dictionary from the environment.

_extra¶

Extra fields for extensibility.

Type:: Dict[str, Any]

ts_env_sent_ns¶

Timestamp (ns) when this EnvRet was produced.

Type:: int

compute_sent_latency(now_ns: int | None = None) → float[source]¶

Compute latency in seconds from env-sent timestamp to now.

Use for env->policy or env->buffer transfer time (e.g. in policy _rollout_hook or in buffer add_transition).

Parameters:: now_ns – Current time in nanoseconds. If None, uses time.time_ns().
Returns:: Latency in seconds, or 0.0 if ts_env_sent_ns is missing or invalid.

cpu() → EnvRet[source]¶

Move all tensor attributes to CPU.

Returns:: Self with tensors moved to CPU.

cuda(device: str | int | None = None) → EnvRet[source]¶

Move all tensor attributes to CUDA.

Parameters:: device – CUDA device index or string. Defaults to ‘cuda’.
Returns:: Self with tensors moved to CUDA.

env_id: str¶: Environment unique identifier

classmethod fields() → Tuple[source]¶

Get field names excluding internal fields.

Returns:: Tuple of field names for serialization.

classmethod get_defaults() → Dict[str, Any][source]¶

Get default values for all serializable fields.

Returns:: Dictionary mapping field names to default values.

info: Dict[str, Any]¶: Additional info from the environment

last_reward: float = 0.0¶: Reward received after the last step

last_terminated: bool = False¶: Whether the episode has terminated after the last step

last_truncated: bool = False¶: Whether the episode has been truncated after the last step

mark_env_sent() → EnvRet[source]¶

Record the timestamp when the env return is sent.

Returns:: Self for method chaining.

numpy() → EnvRet[source]¶

Convert all tensor attributes to numpy arrays.

Returns:: Self with tensors converted to numpy arrays.

observation: Any¶: Observation after environment step/reset

to_dict() → Dict[str, Any][source]¶

Convert EnvRet to dictionary.

Excludes ‘env_id’ and ‘ts_env_sent_ns’. Includes _extra fields if present.

Returns:: Dictionary representation of the environment return.
Raises:: KeyError – If _extra contains keys conflicting with existing fields.

ts_env_sent_ns: int¶: Timestamp (ns) when this EnvRet is produced and sent by the env actor

class rlightning.types.MultiAgentEnvRet(env_id: str, observation: Any, last_reward: Dict[str, float]=0.0, last_terminated: Dict[str, bool]=False, last_truncated: Dict[str, bool]=False, info: Dict[str, ~typing.Any]=<factory>, _extra: Dict[str, ~typing.Any]=<factory>, ts_env_sent_ns: int = <factory>)[source]¶

Bases: EnvRet

Environment return for multi-agent interactions.

Extends EnvRet with dictionary-based rewards and termination flags for multiple agents.

last_reward¶

Dictionary mapping agent IDs to rewards.

Type:: Dict[str, float]

last_terminated¶

Dictionary mapping agent IDs to termination flags.

Type:: Dict[str, bool]

last_truncated¶

Dictionary mapping agent IDs to truncation flags.

Type:: Dict[str, bool]

class rlightning.types.MultiAgentPolicyResponse(env_id: str, **data: Any)[source]¶

Bases: PolicyResponse

Policy response for multi-agent interactions.

Extends PolicyResponse with support for multiple agents, where actions are stored in a dictionary keyed by agent ID.

static make_example(action_spaces: Dict[str, MockModule('gymnasium.Space')], env_id: str | None = None, **data: Any) → MultiAgentPolicyResponse[source]¶

Create an example MultiAgentPolicyResponse.

Parameters:

action_spaces – Dictionary mapping agent IDs to action spaces.
env_id – Optional environment identifier.
**data – Additional response data.

Returns:

MultiAgentPolicyResponse with sampled actions for all agents.

class rlightning.types.PolicyResponse(env_id: str, **data: Any)[source]¶

Bases: SimpleNamespace

Policy response container for single-agent interactions.

Holds the action and any additional data produced by a policy for a single environment step.

env_id¶: Environment identifier this response is for.

Additional attributes are set dynamically via **data.

compute_sent_latency(now_ns: int | None = None) → float[source]¶

Compute latency in seconds from env-sent timestamp to now.

Use for env->policy or env->buffer transfer time (e.g. in policy _rollout_hook or in buffer add_transition).

Parameters:: now_ns – Current time in nanoseconds. If None, uses time.time_ns().
Returns:: Latency in seconds, or 0.0 if ts_env_sent_ns is missing or invalid.

cpu() → PolicyResponse[source]¶

Move all tensor attributes to CPU.