Parameterized OmSh Actions¶

OmSh supports three augmented-action variants in addition to the existing OmSh-Egalitarian-Pure control:

The three new modes use ParameterizedTransitionShield with ParameterizedIPPO. They are named for exports in notebooks.utils.shielding.PARAMETERIZED_OMSH_VARIANTS.

All variants share the opponent-uncertainty settings of the ordinary shield. The default is the robust monotone_floor constraint with Bayesian expected reward ranking. all_levels keeps every per-level constraint; the bayesian_mixture ablation projects against the single posterior-weighted successor distribution while retaining the all-level continuation certificate.

On-policy boundary¶

The focal Gymnasium action space is a Dict containing:

• categorical: one primitive-action index for the pure variant or two edge endpoint indices for a mixed variant
• continuous: raw Gaussian samples containing the edge coordinate, learned successor-budget proposals, or both

The policy stores and scores the exact raw Gaussian samples z. The wrapper applies sigmoid(z) as part of the environment transition, so PPO does not use a transformed-density correction. State-dependent continuous masks exclude unused padded successor slots from sampling, log probability, entropy, and gradient. The primitive action subsequently sampled by the wrapper is not the PPO action and is not substituted into the rollout record.

ParameterizedPolicyAdapter retains both factors after training. Its deterministic action uses categorical modes and Gaussian means; seeded actions sample both factors. Like the ordinary IPPO export, it is unmasked because an inference adapter does not own the live environment state. Safety remains the wrapper's responsibility.

Pairwise action projection¶

For policy action (i, j, z_t), let x = sigmoid(z_t) and

v(x) = (1 - x) e_i + x e_j.

For every active robust opponent level l, the wrapper intersects

c[l, i] + x (c[l, j] - c[l, i]) <= q,
c[l, a] = p[l, a]^T beta_f,

with [0, 1], then clips x to the resulting interval. This is an exact one-dimensional projection across all active robust constraints. It can retain a feasible mixture even when its two endpoints violate different level constraints.

If the selected edge has no feasible point, the runtime supplies local.safe_action as the fallback. This is the existing reward-aware OmSh replacement: maximize the configured reward aggregate over budget-admissible pure actions, then break ties by lower aggregated risk and action index. The reward aggregate is posterior expectation by default; minimum reward across the considered levels remains an ablation. The standalone projection primitive defaults to the minimum-risk admissible pure action when no explicit fallback is supplied.

After projection, the wrapper forms each level's successor distribution

p[l, v] = sum_a v[a] p[l, a]

and samples the realized primitive action from v. Pairwise mixtures strictly contain pure actions (i == j, or an edge endpoint), but they do not represent every vertex of a robust action polytope. With L independent active safety inequalities, a vertex can require support up to min(|A|, L + 1).

Learned successor-budget projection¶

For the current successor union and certificate beta_f, the Gaussian samples propose

W_tilde[j] = beta_f[j] + (1 - beta_f[j]) sigmoid(z[j]),
d = W_tilde - beta_f.

For the pure or already-projected mixed successor distributions p[l], the wrapper computes

lambda = clip_[0,1] min(
    1,
    min_l (q - p[l]^T beta_f) / (p[l]^T d)
)
W = beta_f + lambda d.

Denominators at or below the projection tolerance are skipped, and an empty inner minimum gives lambda = 1. The implementation explicitly verifies beta_f <= W <= 1 and p[l]^T W <= q + eps for every active level. A numerical failure falls back to W = beta_f; an infeasible base certificate raises.

Learning W is reward-aware through the augmented MDP: the realized budget changes the next state's feasible action set, and PPO attributes future return to the exact proposal that produced it. Projection keeps safety hard rather than adding reward to the constraint objective. This permits useful budget allocation but does not guarantee that PPO finds a globally optimal allocation; many-to-one projections can introduce flat regions and variance.

The combined mode first projects v under beta_f, then projects W using the induced p[l, v]. This conservative composition avoids a joint nonlinear projection.

Runtime invariants¶

• Missing state/action coverage is represented by the explicit unsafe sink with certificate 1.0.
• The realized next state receives exactly its component of projected W; an unlisted next state receives budget 1.0.
• The floor certificate is robust across all levels at or above the current monotone floor. Raising the floor therefore preserves next-state feasibility.
• Every emitted action distribution is finite, nonnegative, and normalized.
• The mixed wrapper's categorical endpoint heads remain unmasked. Masking endpoints independently would incorrectly remove mixtures whose endpoints violate different constraints. The pure learned-budget head retains the current admissible-action mask.

Experiment wiring¶

The curated environment notebooks expose shield_action_mode, defaulting to pure. Selecting a parameterized mode makes the shared shield factory return ParameterizedTransitionShield, makes run_jax_stage select ParameterizedIPPO, and adds the action mode to algorithm, timing, history, checkpoint, diagnostic, and plot names. The canonical pure mode uses the unsuffixed baseline artifact names.

For example, a tagged mixed-action tuning run can use:

uv run marimo export ipynb notebooks/matrix/inspection/experiments_mo.py \
  -- \
  run_ippo_shielded=True \
  shield_action_mode=egalitarian_mixed \
  experiment_tag=tune_inspection_mixed_v1 \
  seed_offset=100

The lower-level factory remains available for custom experiment drivers:

from notebooks.utils.shielding import (
    PARAMETERIZED_OMSH_VARIANTS,
    make_parameterized_shielded_env_factory,
)
from src.rl import PARAMETERIZED_IPPO_DEFAULT_CONFIG, ParameterizedIPPO

mode = PARAMETERIZED_OMSH_VARIANTS["OmSh-LearnedBudget-Mixed"]
create_parameterized_env = make_parameterized_shielded_env_factory(
    env_factory=create_env,
    bundle_result=shield_bundle_result,
    max_risk=config["max_risk"],
    action_mode=mode,
    device=device,
    opponent_mode=config["shield_opponent_mode"],
    reward_mode=config["shield_reward_mode"],
)

All four conditions start PPO from fresh initialization; the WM/OM pretraining artifacts remain shared, just as they are for the existing shielded control.

Useful update diagnostics include the mean radial projection scale, mean mixed support, edge-fallback rate, and the ordinary PPO loss/entropy/KL metrics. The environment info also exposes proposed/applied edge coordinates, feasible intervals, action distributions, proposed/projected budgets, and numerical fallback flags according to the configured shield debug level.

Relationship to the paper and MASA¶

Learning successor budgets follows the encoded-action construction in the probabilistic-shielding paper. MASA v1 also learns budget proposals, while MASA v2 uses an egalitarian fill. This implementation does not copy MASA's action mapping: v1 may average multiple intersections and have broader support, and v2 can mix a pair-supported base distribution with a third safest action. OmSh instead guarantees pairwise support and handles all active opponent-level half-spaces explicitly.

Full-simplex action projection is deliberately deferred until pairwise experiments show a remaining reward gap.