marla.learning

The recurrent PPO policy, its sub-modules, rollout collection, and the PPO update itself.

Policy and its components

Ties the graph encoder, action encoder, GRU, base scorer, and critic together into a single live rollout-collection step.

This module handles exactly one environment instance at a time (MARLA does not vectorize environments – see spec section 25 non-goals), so all “batch” dimensions here are size 1. Fixed-size, padded multi-sequence batching for PPO minibatch replay is built on the same components in Milestone 4’s learning/ppo.py.

class marla.learning.recurrent_policy.PolicyStepOutput(z: 'Tensor', base_logits: 'Tensor', base_probs: 'Tensor', action_embeddings: 'Tensor', value: 'Tensor')[source]

Bases: object

Parameters:

  • z (torch.Tensor)

  • base_logits (torch.Tensor)

  • base_probs (torch.Tensor)

  • action_embeddings (torch.Tensor)

  • value (torch.Tensor)

action_embeddings: torch.Tensor
base_logits: torch.Tensor
base_probs: torch.Tensor
value: torch.Tensor
z: torch.Tensor
class marla.learning.recurrent_policy.RecurrentPolicy(*args, **kwargs)[source]

Bases: Module

The trainable RL Orchestrator policy (spec section 11).

Only this module’s parameters are optimized by PPO – the Plan Maker is frozen and lives entirely outside this class.

Parameters:
advance_recurrent_state(step_output, selected_action_index, training_reward, query)[source]

Build the next step’s RecurrentState after an action is selected.

Parameters:
Return type:

RecurrentState

apply_advice(z, base_logits, confidence)[source]

Accepted-advice residual adjustment (spec section 15).

confidence must already be in the same order as base_logits (i.e. the legal-action order), reconstructed by stable action ID – never by raw vector position.

Parameters:

  • z (torch.Tensor)

  • base_logits (torch.Tensor)

  • confidence (torch.Tensor)

Return type:

AdvisedOutput

property device: torch.device
initial_recurrent_state()[source]

The recurrent state at the start of a new episode.

Must be used at every episode boundary – hidden state is never propagated across episodes (spec section 17).

Return type:

RecurrentState

step(graph_observation, legal_actions, recurrent_state)[source]
Parameters:
Return type:

PolicyStepOutput

class marla.learning.recurrent_policy.RecurrentState(z, previous_action_embedding, previous_reward=0.0, previous_query=0.0)[source]

Bases: object

The GRU hidden state plus everything needed to build the next x_t.

Parameters:

  • z (torch.Tensor)

  • previous_action_embedding (torch.Tensor)

  • previous_reward (float)

  • previous_query (float)

previous_action_embedding: torch.Tensor
previous_query: float = 0.0
previous_reward: float = 0.0
z: torch.Tensor

Torch Geometric GraphSAGE encoder (spec section 12).

Produces per-node embeddings (used as target-host embeddings for the action encoder) and a global mean-pooled graph embedding (fed into the GRU).

class marla.learning.graph_encoder.GraphEncoder(*args, **kwargs)[source]

Bases: Module

Parameters:

  • input_dim (int)

  • hidden_dim (int)

  • layers (int)

forward(data)[source]

Returns (node_embeddings, graph_embeddings).

data.batch is required when data holds more than one graph (a Batch); for a single graph it may be omitted (global_mean_pool treats all nodes as one graph).

Parameters:

data (torch_geometric.data.Batch)

Return type:

tuple[torch.Tensor, torch.Tensor]

Per-action embedding \(e_{t,i} = f_{action}(E_{type}, h_{target}, E_{parameters})\).

Non-target actions (finish) use a learned no-target embedding in place of a target-host embedding (spec section 12). Action-specific parameters (exploit/privesc service/os/process names) are summarized as a small fixed-size presence indicator rather than embedded by name, so the model does not bake in a per-scenario service/OS vocabulary – scenarios can introduce new service/OS/process names without changing the model shape.

class marla.learning.action_encoder.ActionEncoder(*args, **kwargs)[source]

Bases: Module

Parameters:

  • node_embedding_size (int)

  • action_type_embedding_size (int)

  • hidden_size (int)

encode_descriptors(descriptors, node_embeddings, node_key_to_index, device)[source]

Convenience wrapper: build encoder inputs directly from descriptors.

Parameters:
Return type:

torch.Tensor

forward(type_ids, target_embeddings, has_target_mask, parameter_feats)[source]

All inputs are batched over actions: shapes (N, ...).

target_embeddings rows are ignored (replaced by the learned no-target embedding) wherever has_target_mask is False.

Parameters:

  • type_ids (torch.Tensor)

  • target_embeddings (torch.Tensor)

  • has_target_mask (torch.Tensor)

  • parameter_feats (torch.Tensor)

Return type:

torch.Tensor

marla.learning.action_encoder.action_type_id(action_type)[source]
Parameters:

action_type (str)

Return type:

int

marla.learning.action_encoder.parameter_features(parameters)[source]
Parameters:

parameters (dict[str, object])

Return type:

torch.Tensor

GRU recurrent state (spec section 13).

z_t = GRUCell(x_t, z_{t-1}) with x_t = [g_t, E(a_{t-1}), r~_{t-1}, q_{t-1}]. At episode start, z is zero, the previous action is a learned start token, and the previous reward/query are zero. All tensors are batch-first (B = 1 for single-instance live rollout collection, B > 1 for PPO minibatch replay in Milestone 4) so the same module serves both.

class marla.learning.recurrent_core.RecurrentCore(*args, **kwargs)[source]

Bases: Module

Parameters:

  • graph_embedding_size (int)

  • action_embedding_size (int)

  • hidden_size (int)

forward(graph_embedding, previous_action_embedding, previous_reward, previous_query, previous_hidden)[source]
Parameters:

  • graph_embedding (torch.Tensor)

  • previous_action_embedding (torch.Tensor)

  • previous_reward (torch.Tensor)

  • previous_query (torch.Tensor)

  • previous_hidden (torch.Tensor)

Return type:

torch.Tensor

initial_hidden_state(batch_size, device)[source]
Parameters:

  • batch_size (int)

  • device (torch.device)

Return type:

torch.Tensor

initial_previous_action_embedding(batch_size, device)[source]
Parameters:

  • batch_size (int)

  • device (torch.device)

Return type:

torch.Tensor

Dynamic candidate-action scorer (spec section 13).

b_{t,i} = w^T tanh(W_z z_t + W_e e_{t,i}). Candidate sets vary in size per sample, so actions are passed flattened across the batch together with an action_to_sample index vector – the same pattern Torch Geometric uses for variable-size graphs in a Batch (see data.batch).

class marla.learning.base_scorer.BaseActionScorer(*args, **kwargs)[source]

Bases: Module

Parameters:

  • recurrent_hidden_size (int)

  • action_hidden_size (int)

  • scorer_hidden_size (int)

forward(z, action_embeddings, action_to_sample)[source]

Returns raw (unnormalized) logits, one per action row.

z: (B, recurrent_hidden_size). action_embeddings: (N_total, action_hidden_size). action_to_sample: (N_total,) int64, values in [0, B).

Parameters:

  • z (torch.Tensor)

  • action_embeddings (torch.Tensor)

  • action_to_sample (torch.Tensor)

Return type:

torch.Tensor

Value critic: a linear head over the GRU recurrent state.

class marla.learning.critic.Critic(*args, **kwargs)[source]

Bases: Module

Parameters:

hidden_size (int)

forward(z)[source]

z: (B, hidden_size) -> (B,).

Parameters:

z (torch.Tensor)

Return type:

torch.Tensor

Learned binary query gate (spec section 14).

The query decision happens before the Plan Maker is consulted: it is a function of the base policy alone (recurrent state, base entropy, top-two margin, candidate-set size, and the configured consultation cost).

class marla.learning.query_gate.QueryGate(*args, **kwargs)[source]

Bases: Module

p_t^q = sigma(f_q[z_t, H_t^0, Delta_t^0, N_t, kappa]).

Parameters:

  • recurrent_hidden_size (int)

  • mlp_hidden_size (int)

forward(z, entropy, margin, num_actions, kappa)[source]

All of entropy/margin/num_actions/kappa are 0-dim tensors.

Parameters:

  • z (torch.Tensor)

  • entropy (torch.Tensor)

  • margin (torch.Tensor)

  • num_actions (torch.Tensor)

  • kappa (torch.Tensor)

Return type:

torch.Tensor

marla.learning.query_gate.compute_base_entropy(base_probs)[source]

H_t^0 = -sum_i pi^0(a_i) log pi^0(a_i).

Parameters:

base_probs (torch.Tensor)

Return type:

torch.Tensor

marla.learning.query_gate.compute_top_two_margin(base_logits)[source]

Delta_t^0 = b_(1) - b_(2); zero when only one legal action exists.

Parameters:

base_logits (torch.Tensor)

Return type:

torch.Tensor

Advice processing and learned trust (spec section 15).

Confidence scores from the Plan Maker are not assumed calibrated probabilities: they are clipped, converted to log-odds, and z-score normalized before being used as a residual adjustment to the base logits, scaled by a learned per-step trust coefficient and a single learned global scale.

class marla.learning.advice.AdviceScale(*args, **kwargs)[source]

Bases: Module

alpha = softplus(alpha_bar); a single learned global scalar (not per-step).

forward()[source]
Return type:

torch.Tensor

class marla.learning.advice.AdviceSummary(mean: 'Tensor', std: 'Tensor', max: 'Tensor', top1_minus_top2: 'Tensor', entropy: 'Tensor')[source]

Bases: object

Parameters:

  • mean (torch.Tensor)

  • std (torch.Tensor)

  • max (torch.Tensor)

  • top1_minus_top2 (torch.Tensor)

  • entropy (torch.Tensor)

entropy: torch.Tensor
max: torch.Tensor
mean: torch.Tensor
std: torch.Tensor
top1_minus_top2: torch.Tensor
class marla.learning.advice.AdvisedOutput(final_logits: 'Tensor', beta: 'Tensor', alpha: 'Tensor', normalized_advice: 'Tensor')[source]

Bases: object

Parameters:

  • final_logits (torch.Tensor)

  • beta (torch.Tensor)

  • alpha (torch.Tensor)

  • normalized_advice (torch.Tensor)

alpha: torch.Tensor
beta: torch.Tensor
final_logits: torch.Tensor
normalized_advice: torch.Tensor
class marla.learning.advice.TrustHead(*args, **kwargs)[source]

Bases: Module

beta_t = sigma(f_beta[z_t, S(c_t), A_t]).

Parameters:

  • recurrent_hidden_size (int)

  • mlp_hidden_size (int)

AGREEMENT_DIM = 2
SUMMARY_DIM = 5
forward(z, summary, agreement)[source]
Parameters:

  • z (torch.Tensor)

  • summary (torch.Tensor)

  • agreement (torch.Tensor)

Return type:

torch.Tensor

marla.learning.advice.clip_and_logit(confidence)[source]

c_i -> u_i = logit(clip(c_i, eps, 1-eps)).

Parameters:

confidence (torch.Tensor)

Return type:

torch.Tensor

marla.learning.advice.compute_advice_summary(log_odds)[source]

S(c) = [mean, std, max, top1-top2, H(softmax(u))].

Parameters:

log_odds (torch.Tensor)

Return type:

AdviceSummary

marla.learning.advice.compute_agreement_features(base_logits, advice_log_odds)[source]

[top-action agreement (0/1), correlation between base logits and PM evidence].

Correlation is 0 when undefined (zero variance in either vector) – spec section 15: “If correlation is undefined, use zero.”

Fixed, non-differentiable features for the trust head: the Plan Maker’s output is external data, not a trainable parameter, so there is no gradient PPO could usefully take through this comparison – base_logits keeps its own graph connection for the outer residual sum in RecurrentPolicy.apply_advice(), this is a separate detached copy.

Parameters:

  • base_logits (torch.Tensor)

  • advice_log_odds (torch.Tensor)

Return type:

torch.Tensor

marla.learning.advice.normalize_advice(log_odds)[source]

z-score normalize log-odds; all-zero when variance is ~0 (spec section 15).

Parameters:

log_odds (torch.Tensor)

Return type:

torch.Tensor

marla.learning.advice.summary_to_tensor(summary)[source]
Parameters:

summary (AdviceSummary)

Return type:

torch.Tensor

Compound-decision math shared by live rollout collection and PPO replay.

Both call sites need exactly the same computation: at collection time to record the “old” probabilities, and at replay time to recompute “new” probabilities under updated parameters for the PPO ratio. Sharing this module is what makes that guarantee structural rather than a matter of remembering to keep two implementations in sync.

Replay never samples: it reuses the stored sampled_query flag, the stored selected_action_index, and (critically) the exact stored Plan Maker confidence vector – never recomputing or re-fetching any of them.

class marla.learning.decision.FinalDecision(final_logits: 'Tensor', beta: 'Tensor | None', alpha: 'Tensor | None', normalized_advice: 'Tensor | None')[source]

Bases: object

Parameters:

  • final_logits (Tensor)

  • beta (Tensor | None)

  • alpha (Tensor | None)

  • normalized_advice (Tensor | None)

alpha: Tensor | None
beta: Tensor | None
final_logits: Tensor
normalized_advice: Tensor | None
class marla.learning.decision.JointLogProbability(joint_log_prob: 'Tensor', query_log_prob: 'Tensor', action_log_prob: 'Tensor')[source]

Bases: object

Parameters:

  • joint_log_prob (torch.Tensor)

  • query_log_prob (torch.Tensor)

  • action_log_prob (torch.Tensor)

action_log_prob: torch.Tensor
joint_log_prob: torch.Tensor
query_log_prob: torch.Tensor
marla.learning.decision.compute_final_decision(policy, step_output, sampled_query, plan_maker_confidence)[source]

Base logits, or the advised residual adjustment when queried and accepted.

Rejected advice forces beta=0 (spec section 15), which makes the resulting distribution numerically identical to the base policy while still reporting a real (if unused) alpha for metrics.

Parameters:
Return type:

FinalDecision

marla.learning.decision.compute_joint_log_probability(query_probability, sampled_query, final_logits, selected_action_index)[source]

log p_t = log pi^q(q_t|z_t) + (1-q_t) log pi^0(a_t) + q_t log pi^PM(a_t) (spec section 16).

The single final_logits softmax already is pi^0 when not queried (or when advice was rejected, since beta=0) and pi^PM when queried and accepted, so one formula covers both branches of the spec’s piecewise definition without a special case.

Parameters:

  • query_probability (torch.Tensor)

  • sampled_query (bool)

  • final_logits (torch.Tensor)

  • selected_action_index (int)

Return type:

JointLogProbability

marla.learning.decision.compute_query_probability(policy, step_output, num_actions, kappa)[source]

p_t^q, before sampling (spec section 14).

Parameters:
Return type:

torch.Tensor

Rollout, training, and optimization

Recurrent rollout collection (spec sections 16-17).

Runs the compound-decision loop against a live NasimEmuAdapter and records everything PPO needs to replay each transition without ever calling the environment or the Plan Maker again. Baseline instances are constructed with consultation_enabled=False and never touch the query-gate/advice machinery at all; assisted instances receive a consult_fn callback that performs the actual Gatekeeper round-trip (kept out of this module so it stays decoupled from SPADE – see marla.agents.advisory_client).

class marla.learning.rollout.ConsultationResult(status, scores, request_id, latency_ms=None)[source]

Bases: object

What a consult_fn callback returns; decoupled from SPADE/Gatekeeper types.

Parameters:
latency_ms: float | None = None
request_id: str
scores: dict[str, float] | None
status: str
class marla.learning.rollout.EpisodeSummary(run_id: 'str', episode_id: 'int', seed: 'int', goal_success: 'bool', nasimemu_return: 'float', training_return: 'float', environment_steps: 'int', steps_to_goal: 'int | None', episode_seconds: 'float', finish_reason: 'str', consultation_count: 'int' = 0, consultation_cost: 'float' = 0.0, schema_rejection_count: 'int' = 0, rollout: 'int | None' = None, is_eval: 'bool' = False)[source]

Bases: object

Parameters:

  • run_id (str)

  • episode_id (int)

  • seed (int)

  • goal_success (bool)

  • nasimemu_return (float)

  • training_return (float)

  • environment_steps (int)

  • steps_to_goal (int | None)

  • episode_seconds (float)

  • finish_reason (str)

  • consultation_count (int)

  • consultation_cost (float)

  • schema_rejection_count (int)

  • rollout (int | None)

  • is_eval (bool)

consultation_cost: float = 0.0
consultation_count: int = 0
environment_steps: int
episode_id: int
episode_seconds: float
finish_reason: str
goal_success: bool
is_eval: bool = False
nasimemu_return: float
rollout: int | None = None
run_id: str
schema_rejection_count: int = 0
seed: int
steps_to_goal: int | None
training_return: float
class marla.learning.rollout.RolloutCollector(adapter, policy, run_id, base_seed, consultation_enabled=False, consultation_cost=0.0, consult_fn=None, stop_event=None, deterministic=False)[source]

Bases: object

Collects fixed-length rollout windows, resuming across calls.

One NasimEmuAdapter/RecurrentPolicy pair per instance, matching MARLA’s non-vectorized, single-environment design.

Parameters:
async collect(num_steps)[source]
Parameters:

num_steps (int)

Return type:

tuple[list[StepRecord], list[EpisodeSummary]]

class marla.learning.rollout.StepRecord(run_id, episode_id, environment_step, observation_id, graph_data, node_key_to_index, legal_action_descriptors, initial_gru_hidden_state, previous_action_embedding, previous_training_reward, previous_query, base_logits, final_logits, selected_action_index, old_action_log_probability, old_joint_log_probability, critic_value, nasimemu_reward, consultation_cost, training_reward, terminated, truncated, bootstrap_value=None, sampled_query=False, old_query_probability=None, old_query_log_probability=None, plan_maker_scores_in_action_order=None, plan_maker_validation_status=None, plan_maker_response_status=None, plan_maker_request_id=None, plan_maker_latency_ms=None, plan_maker_artifact_path=None, normalized_advice=None, beta=None, alpha=None)[source]

Bases: object

One compound decision, with everything needed to replay it under PPO.

Parameters:

  • run_id (str)

  • episode_id (int)

  • environment_step (int)

  • observation_id (str)

  • graph_data (torch_geometric.data.Data)

  • node_key_to_index (dict[str, int])

  • legal_action_descriptors (list[ActionDescriptor])

  • initial_gru_hidden_state (torch.Tensor)

  • previous_action_embedding (torch.Tensor)

  • previous_training_reward (float)

  • previous_query (bool)

  • base_logits (torch.Tensor)

  • final_logits (torch.Tensor)

  • selected_action_index (int)

  • old_action_log_probability (float)

  • old_joint_log_probability (float)

  • critic_value (float)

  • nasimemu_reward (float)

  • consultation_cost (float)

  • training_reward (float)

  • terminated (bool)

  • truncated (bool)

  • bootstrap_value (float | None)

  • sampled_query (bool)

  • old_query_probability (float | None)

  • old_query_log_probability (float | None)

  • plan_maker_scores_in_action_order (list[float] | None)

  • plan_maker_validation_status (str | None)

  • plan_maker_response_status (str | None)

  • plan_maker_request_id (str | None)

  • plan_maker_latency_ms (float | None)

  • plan_maker_artifact_path (str | None)

  • normalized_advice (list[float] | None)

  • beta (float | None)

  • alpha (float | None)

alpha: float | None = None
base_logits: torch.Tensor
beta: float | None = None
bootstrap_value: float | None = None
consultation_cost: float
critic_value: float
environment_step: int
episode_id: int
final_logits: torch.Tensor
graph_data: torch_geometric.data.Data
initial_gru_hidden_state: torch.Tensor
legal_action_descriptors: list[ActionDescriptor]
nasimemu_reward: float
node_key_to_index: dict[str, int]
normalized_advice: list[float] | None = None
observation_id: str
old_action_log_probability: float
old_joint_log_probability: float
old_query_log_probability: float | None = None
old_query_probability: float | None = None
plan_maker_artifact_path: str | None = None
plan_maker_latency_ms: float | None = None
plan_maker_request_id: str | None = None
plan_maker_response_status: str | None = None
plan_maker_scores_in_action_order: list[float] | None = None
plan_maker_validation_status: str | None = None
previous_action_embedding: torch.Tensor
previous_query: bool
previous_training_reward: float
run_id: str
sampled_query: bool = False
selected_action_index: int
terminated: bool
training_reward: float
truncated: bool
async marla.learning.rollout.run_evaluation_episodes(policy, adapter, run_id, num_episodes, seed_start, consultation_enabled=False, consultation_cost=0.0, consult_fn=None)[source]

Runs num_episodes deterministic (greedy) episodes with the current policy weights and reports their returns, contributing nothing to the PPO buffer – an EvalCallback-style periodic evaluation pass (see e.g. stable-baselines3), not a genuine held-out generalization test: MARLA’s config has a single environment.scenario, not a train/test scenario split, so this measures the current policy’s performance without exploration noise on the same scenario training uses, not generalization to unseen scenarios.

Reuses adapter rather than building a second instance: training and evaluation never run concurrently (this is awaited strictly between rollout-collection phases in run_training_loop), and NasimEmuAdapter holds no state of its own between reset() calls.

Parameters:
Return type:

list[EpisodeSummary]

Recurrent PPO training loop, shared by the baseline and assisted variants.

Ties together RolloutCollector, compute_gae(), and optimize() into repeated rollout-collect/PPO-update cycles. Native async def throughout: the assisted variant’s per-step consultation is a real await on the Gatekeeper round-trip, and running natively on the event loop (rather than in a background thread) is what lets other local-mode agents’ behaviours keep being serviced during training. The SPADE RL Orchestrator (Milestone 5) awaits this directly; this module is also directly usable standalone for the baseline variant and for tests.

class marla.learning.trainer.TrainingResult(policy: 'RecurrentPolicy', optimizer: 'torch.optim.Optimizer', episode_summaries: 'list[EpisodeSummary]' = <factory>, update_metrics: 'list[dict[str, float]]' = <factory>, environment_steps: 'int' = 0, all_records: 'list[StepRecord]' = <factory>, stopped_by_user: 'bool' = False, eval_episode_summaries: 'list[EpisodeSummary]' = <factory>)[source]

Bases: object

Parameters:
all_records: list[StepRecord]
environment_steps: int = 0
episode_summaries: list[EpisodeSummary]
eval_episode_summaries: list[EpisodeSummary]
optimizer: torch.optim.Optimizer
policy: RecurrentPolicy
stopped_by_user: bool = False
update_metrics: list[dict[str, float]]
marla.learning.trainer.build_policy_and_optimizer(config, device, consultation_enabled=False)[source]
Parameters:

  • config (Config)

  • device (torch.device)

  • consultation_enabled (bool)

Return type:

tuple[RecurrentPolicy, torch.optim.Optimizer]

async marla.learning.trainer.run_baseline_training(config, scenario_path, num_rollouts, device=None, seed=None)[source]

Run num_rollouts collect/optimize cycles of baseline recurrent PPO.

Parameters:

  • config (Config)

  • scenario_path (str)

  • num_rollouts (int)

  • device (torch.device | None)

  • seed (int | None)

Return type:

TrainingResult

async marla.learning.trainer.run_training_loop(policy, optimizer, adapter, run_id, ppo_config, sequence_length, num_rollouts, device, seed, consultation_enabled=False, consultation_cost=0.0, consult_fn=None, stop_event=None, eval_episodes=0, eval_every_rollouts=1)[source]

Repeated rollout-collect/PPO-update cycles over pre-built components.

Used both by run_baseline_training() (which builds everything itself, for standalone use and tests) and by the SPADE RL Orchestrator, which resolves the device and constructs the policy before connecting to XMPP – a local-model failure must prevent the agent from ever reporting itself ready (spec section 18).

Parameters:
Return type:

TrainingResult

Generalized Advantage Estimation with correct episode-boundary handling.

Three cases determine the “next value” used at each step:

  • terminated: a true terminal transition (FINISH). The next value is 0 – there is no future return to bootstrap.

  • truncated: the episode was cut short (max_episode_steps). The caller must supply a bootstrap_values[t] estimate of the state that followed, since the next buffer row (if any) belongs to a fresh episode with a reset GRU state and is not a valid continuation.

  • Neither, but t is the last row of the rollout window: the PPO horizon ended mid-episode. The caller must again supply bootstrap_values[t] (the value of continuing under the current policy).

In every other case, the next value is simply values[t + 1] – the same episode continues into the next buffer row.

marla.learning.gae.compute_gae(rewards, values, terminated, truncated, bootstrap_values, gamma, gae_lambda)[source]

Returns (advantages, returns), both length len(rewards).

Parameters:
Return type:

tuple[list[float], list[float]]

Recurrent PPO optimization (spec sections 16-17).

Sequences are chunked per-episode (never spanning an episode boundary) and replayed by re-running RecurrentPolicy.step() forward through the chunk with gradients enabled, seeded by each chunk’s stored initial hidden state and previous-action embedding (a “stored state” truncated-BPTT scheme: the chunk boundary is a fixed input, not backpropagated through). Because NASimEmu observations are already fixed in the buffer, this recomputes only the encoding of each transition under the current parameters – it never touches the environment or the Plan Maker again. For the assisted variant, the exact stored Plan Maker confidence vector is reused verbatim; the Plan Maker itself is never re-invoked during replay.

Sequences are batched and shuffled as whole units (never individual transitions), and each chunk stays within one episode by construction, so there is no cross-episode hidden-state leakage to mask out. Because MARLA does not vectorize environments, chunks are replayed one at a time rather than as a single padded tensor – correctness first, matching the project’s “NASimEmu simulation correctness is the priority” scoping for v0.1.

class marla.learning.ppo.ReplayOutputs(new_joint_log_probs: 'Tensor', new_values: 'Tensor', action_entropies: 'Tensor', query_entropies: 'Tensor | None')[source]

Bases: object

Parameters:

  • new_joint_log_probs (Tensor)

  • new_values (Tensor)

  • action_entropies (Tensor)

  • query_entropies (Tensor | None)

action_entropies: Tensor
new_joint_log_probs: Tensor
new_values: Tensor
query_entropies: Tensor | None
class marla.learning.ppo.SequenceChunk(records: 'list[StepRecord]', advantages: 'list[float]', returns: 'list[float]')[source]

Bases: object

Parameters:
advantages: list[float]
records: list[StepRecord]
returns: list[float]
marla.learning.ppo.build_sequence_chunks(records, advantages, returns, sequence_length)[source]

Group consecutive same-episode records into chunks of at most sequence_length.

A chunk never spans two episodes: whenever episode_id changes, a new chunk starts, which is exactly what prevents hidden-state leakage across episode boundaries during replay.

Parameters:
Return type:

list[SequenceChunk]

marla.learning.ppo.optimize(policy, optimizer, records, advantages, returns, ppo_config, sequence_length, minibatch_sequences, device, rng, consultation_cost=0.0)[source]

Runs ppo_config.epochs passes of shuffled-minibatch updates over the rollout.

Parameters:
Return type:

list[dict[str, float]]

marla.learning.ppo.ppo_update(policy, optimizer, chunks, ppo_config, device, consultation_cost=0.0)[source]

One gradient step over a minibatch of sequence chunks.

Parameters:
Return type:

dict[str, float]

Masked softmax/entropy utilities for padded, fixed-size PPO minibatches.

Live rollout collection scores a single instance’s legal actions directly (no padding needed). PPO replay (Milestone 4) batches multiple timesteps with different candidate-action counts into fixed-size tensors, padding short rows; padded entries must be masked out of softmax, log-probability, and entropy computations (spec section 13).

marla.learning.masking.masked_entropy(log_probs, mask, dim=-1)[source]
Parameters:

  • log_probs (Tensor)

  • mask (Tensor)

  • dim (int)

Return type:

Tensor

marla.learning.masking.masked_log_softmax(logits, mask, dim=-1)[source]

mask: True = valid/real entry, False = padding.

Parameters:

  • logits (Tensor)

  • mask (Tensor)

  • dim (int)

Return type:

Tensor

marla.learning.masking.masked_softmax(logits, mask, dim=-1)[source]
Parameters:

  • logits (Tensor)

  • mask (Tensor)

  • dim (int)

Return type:

Tensor

Checkpoint save/load for the recurrent PPO policy + optimizer.

class marla.learning.checkpoint.CheckpointMetadata(update_count: 'int', environment_steps: 'int', config_hash: 'str')[source]

Bases: object

Parameters:

  • update_count (int)

  • environment_steps (int)

  • config_hash (str)

config_hash: str
environment_steps: int
update_count: int
marla.learning.checkpoint.load_checkpoint(path, policy, optimizer=None, map_location='cpu')[source]
Parameters:
Return type:

CheckpointMetadata

marla.learning.checkpoint.save_checkpoint(path, policy, optimizer, update_count, environment_steps, config_hash)[source]
Parameters:
Return type:

None