rlightning.types.policy_response

Policy response data structures.

This module provides classes for representing policy action responses, supporting both single-agent and multi-agent scenarios.

class rlightning.types.policy_response.MultiAgentPolicyResponse(env_id: str, **data: Any)

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

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.policy_response.PolicyResponse(env_id: str, **data: Any)

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

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

Move all tensor attributes to CPU.

Returns:

Self with tensors moved to CPU.

cuda(device: str | int | None = None) → PolicyResponse

Move all tensor attributes to CUDA.

Parameters:

device – CUDA device index or string. Defaults to ‘cuda’.

Returns:

Self with tensors moved to CUDA.

static make_example(action_space: MockModule('gymnasium.Space'), env_id: str | None = None, **data: Any) → PolicyResponse

Create an example PolicyResponse.

Parameters:

• action_space – Gymnasium action space for sampling.

• env_id – Optional environment identifier.

• **data – Additional response data.

Returns:

PolicyResponse with sampled action.

mark_policy_sent() → PolicyResponse

Record the timestamp when the policy response is sent.

Returns:

Self for method chaining.

numpy() → PolicyResponse

Convert all tensor attributes to NumPy arrays.

Returns:

Self with tensors converted to NumPy arrays.

to_dict() → Dict[str, Any]

Convert to dictionary excluding env_id.

Returns:

Dictionary of response fields.