rlightning.types.env_rets

Environment return data structures.

This module provides dataclasses for representing environment step/reset returns, supporting both single-agent and multi-agent scenarios.

class rlightning.types.env_rets.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.env_rets.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]