Skip to content

RL API

RL APIs cover baseline learners, shared training data structures, policy adapters, checkpoints, frozen final-policy evaluation, opponent-gap diagnostics, and model export helpers.

Shared Training Helpers

src.rl.base.AgentSpec dataclass 

AgentSpec(
    agent_id: str,
    obs_dim: int,
    action_dim: int,
    action_kind: str,
    action_nvec: tuple[int, ...],
    action_shape: tuple[int, ...],
    action_start: tuple[int, ...],
)

src.rl.base.TrainResult dataclass 

TrainResult(
    history: list[Any],
    metrics: dict[str, list[dict[str, float]]],
    policies: dict[str, "PolicyAdapter"],
    agent_states: dict[str, Any],
)

src.rl.base.EpisodicCostTracker dataclass 

EpisodicCostTracker(
    gamma: float = 1.0,
    initial_return: float = 0.0,
    running_return: float = 0.0,
    running_mass: float = 0.0,
    discount: float = 1.0,
    last_return: float | None = None,
    last_scale: float = 1.0,
)

Track returns from true episode starts across rollout boundaries.

src.rl.base.PolicyAdapter dataclass 

PolicyAdapter(
    agent_id: str,
    params: Any,
    apply_fn: Callable[..., Any],
    obs_dim: int,
    action_kind: str = "discrete",
    action_nvec: tuple[int, ...] = (),
    action_shape: tuple[int, ...] = (),
    action_start: tuple[int, ...] = (),
)

component_action_probs 

component_action_probs(obs: Any) -> tuple[np.ndarray, ...]

Return one categorical distribution per discrete action component.

act 

act(
    obs: Any,
    rng: Any = None,
    action_mask: ndarray | None = None,
) -> int | np.ndarray

Sample (or, without rng, take the mode of) the policy's action.

action_mask is the state-dependent admissible set a shield exposes. On-policy training samples only within it, so a frozen evaluation must too: otherwise the wrapper rewrites inadmissible proposals and the evaluation scores a different controller from the one trained. An all-true mask leaves sampling bit-identical.

src.rl.base.merge_configs 

merge_configs(
    defaults: Mapping[str, Any],
    overrides: Mapping[str, Any] | None,
) -> dict[str, Any]

src.rl.base.infer_agent_specs 

infer_agent_specs(env: ParallelEnv) -> list[AgentSpec]

src.rl.base.compute_gae 

compute_gae(
    rewards: ndarray,
    bootstrap_dones: ndarray,
    episode_ends: ndarray,
    values: ndarray,
    next_values: ndarray,
    gamma: float,
    gae_lambda: float,
) -> tuple[jnp.ndarray, jnp.ndarray]

src.rl.base.normalize_advantages 

normalize_advantages(advantages: ndarray) -> jnp.ndarray

src.rl.base.num_optimizer_steps_for_timesteps 

num_optimizer_steps_for_timesteps(
    timesteps: int,
    rollout_length: int,
    learning_epochs: int,
    requested_mini_batches: int,
) -> int

Return the exact number of optimizer applications for a training run.

Algorithms

src.rl.ippo.IPPO 

IPPO(cfg: dict[str, Any] | None = None, *, seed: int = 0)

src.rl.ippo.ActorCritic

Bases: Module

src.rl.ippo.validate_ippo_config 

validate_ippo_config(
    overrides: dict[str, Any] | None = None,
) -> dict[str, Any]

Merge and validate IPPO configuration, rejecting misspelled options.

src.rl.parameterized_ippo.ParameterizedIPPO 

ParameterizedIPPO(
    cfg: dict[str, Any] | None = None, *, seed: int = 0
)

IPPO over factored categorical/raw-Gaussian augmented shield actions.

src.rl.parameterized_ippo.ParameterizedActorCritic

Bases: Module

src.rl.parameterized_ippo.ParameterizedPolicyAdapter dataclass 

ParameterizedPolicyAdapter(
    agent_id: str,
    params: Any,
    apply_fn: Callable[..., Any],
    obs_dim: int,
    action_kind: str = "discrete",
    action_nvec: tuple[int, ...] = (),
    action_shape: tuple[int, ...] = (),
    action_start: tuple[int, ...] = (),
    continuous_dim: int = 0,
)

Bases: PolicyAdapter

Inference adapter that preserves the complete learned augmented action.

act 

act(
    obs: Any,
    rng: Any = None,
    categorical_masks: ndarray | None = None,
    continuous_mask: ndarray | None = None,
) -> dict[str, np.ndarray]

Sample the augmented action, optionally within the shield's factored masks.

Training samples each categorical head within categorical_masks and zeroes continuous components outside continuous_mask; a frozen evaluation passes the same masks so it scores the controller that was trained.

src.rl.parameterized_ippo.validate_parameterized_ippo_config 

validate_parameterized_ippo_config(
    overrides: dict[str, Any] | None = None,
) -> dict[str, Any]

src.rl.ippo_lag.IPPO_Lagrangian 

IPPO_Lagrangian(cfg: dict[str, Any], *, seed: int = 0)

src.rl.ippo_lag.ActorCriticWithCost

Bases: Module

src.rl.ICPO.ICPO 

ICPO(cfg: dict[str, Any], *, seed: int = 0)

src.rl.ICPO.CPO 

CPO(cfg: dict[str, Any] | None = None, *, seed: int = 0)

Bases: ICPO

Plain CPO preset without ICPO's learned cost shaping extension.

src.rl.ICPO.CategoricalActor 

CategoricalActor(
    obs_dim: int,
    action_dim: int,
    *,
    hidden_sizes: tuple[int, ...],
    activation: str,
)

Bases: Module

src.rl.ICPO.ValueNetwork 

ValueNetwork(
    obs_dim: int,
    *,
    hidden_sizes: tuple[int, ...],
    activation: str,
)

Bases: Module

Checkpoints and Frozen Evaluation

src.rl.checkpoints.save_training_checkpoint 

save_training_checkpoint(
    path: str | Path,
    *,
    agent_states: Mapping[str, Any],
    policies: Mapping[str, Any],
    algorithm: str,
    algorithm_config: Mapping[str, Any],
    trainer_seed: int,
    controller_state: Mapping[str, Any] | None = None,
) -> Path

Write a versioned, restorable final-policy and optimizer checkpoint.

src.rl.checkpoints.load_training_checkpoint 

load_training_checkpoint(
    path: str | Path,
) -> dict[str, Any]

src.rl.checkpoints.load_final_policies 

load_final_policies(
    path: str | Path, *, device: str = "cpu"
) -> dict[str, Any]

Restore final inference adapters from a supported training checkpoint.

src.rl.final_policy_evaluation.evaluate_frozen_policies 

evaluate_frozen_policies(
    *,
    env: Any,
    policies: Mapping[str, Any],
    episodes: int,
    horizons: Sequence[int],
    seed_offset: int = 1000000,
    focal_agent_id: str | None = None,
    confidence: float = 0.95,
    initial_controller_state: Mapping[str, Any]
    | None = None,
) -> dict[str, Any]

Evaluate a frozen joint policy with IID seeded live-environment episodes.

Stateful controllers may expose get_evaluation_state and set_evaluation_state. Their post-training snapshot is restored before every replicate, preventing online opponent-posterior or floor updates in one replicate from changing the next replicate's controller.

src.rl.final_policy_evaluation.validate_live_evaluation_protocol 

validate_live_evaluation_protocol(
    report: Mapping[str, Any], *, stateful_controller: bool
) -> None

Reject live reports that do not use the corrected IID replay protocol.

src.rl.final_policy_evaluation.evaluate_exact_policy_reachability 

evaluate_exact_policy_reachability(
    *,
    graph: Graph,
    policies: Mapping[str, Any],
    agent_ids: Sequence[str],
    unsafe_keys: Sequence[str] | set[str],
    horizons: Sequence[int],
    start_distribution: Mapping[str, float] | None = None,
    tolerance: float = 1e-10,
    max_iterations: int = 1000000,
    mass_tolerance: float = 1e-08,
    roundoff_tolerance: float = 1e-12,
) -> dict[str, Any]

Evaluate finite and eventual unsafe reachability of a stationary joint policy.

src.rl.final_policy_evaluation.write_final_policy_evaluation 

write_final_policy_evaluation(
    report: Mapping[str, Any], path: str | Path
) -> Path

Opponent-Gap Diagnostics

src.rl.opponent_level_gap.masked_action_probs 

masked_action_probs(
    logits: Any, mask: Any | None = None
) -> np.ndarray

Softmax over logits restricted to mask, as the trainers sample.

logits is [B, A]; mask is [B, A] (or None for no mask). An all-true mask gives the plain softmax.

src.rl.opponent_level_gap.joint_action_probs 

joint_action_probs(
    per_opponent: Sequence[ndarray],
) -> np.ndarray

Product law of independent opponents, first opponent most significant.

Matches encode_joint_action: joint index a_1 * A_2 + a_2 for two opponents, and so on.

src.rl.opponent_level_gap.closest_retained_level_tv 

closest_retained_level_tv(
    per_level_tv: ndarray, floors: Any
) -> tuple[np.ndarray, np.ndarray]

Per-sample min_{l >= floor} TV and the level attaining it.

per_level_tv is [M, B]; floors is [B] with entries in [0, M).

src.rl.opponent_level_gap.suffix_min_level_tv 

suffix_min_level_tv(per_level_tv: ndarray) -> np.ndarray

Row f is min_{l >= f} of the per-level TV, for every floor f.

Model Export

src.rl.model_io.save_models 

save_models(
    env: Any,
    models: Mapping[Any, Mapping[str, Module]],
    folder: str | Path,
) -> None