Skip to content

Discrete ICPO Port

OmniSafe: https://github.com/PKU-Alignment/omnisafe

src/rl/ICPO.py now contains the active Independent Constrained Policy Optimisation benchmark used by the experiment notebooks.

Why this is local instead of a thin OmniSafe wrapper:

  • OmniSafe's current CPO stack assumes gymnasium.spaces.Box actions.
  • The benchmark environments in omsh are discrete PettingZoo environments.
  • To keep ICPO as a like-for-like benchmark on the existing suites, the repo ports the core CPO update logic locally and applies it to categorical policies.

Implementation choices:

  • Each agent owns its own actor, reward critic, cost critic, and learned failure predictor.
  • Actors are updated with OmniSafe-style CPO math: conjugate gradient, Fisher-vector products from D_KL(old_policy || new_policy), trust-region case selection, and backtracking line search. The same old-to-new KL direction is used by the Fisher geometry and the sampled line-search check.
  • When the current policy is feasible, line search checks sampled surrogate feasibility as constraint_violation + cost_diff <= 0, matching the single-constraint approximation c + b^T step <= 0. When the current policy is already infeasible, no step inside a small trust region may be able to restore full feasibility; in that case a proposal is accepted only if it strictly reduces the sampled cost surrogate.
  • The analytic step uses the reward-maximizing CPO subproblem: max g^T step subject to c + b^T step <= 0 and step^T H step <= 2 * target_kl. If the unconstrained TRPO step violates the active cost plane but the plane intersects the trust region, ICPO moves to the plane and spends the remaining trust-region radius along the reward-improving tangent direction.
  • Critics are updated with plain MSE regression on reward and shaped cost returns.
  • Raw costs are consumed from infos[agent_id][cost_info_key], defaulting to infos[agent_id]["cost"] for compatibility with the repo's other constrained baselines. If an agent has no cost entry, the trainer uses 0.0.
  • By default the CPO constrained signal is the paper-style shaped upper-bound cost C+ = C + alpha * Delta. Delta is the per-agent failure predictor's probability that a raw cost will occur within cost_shaping_horizon rollout steps after (obs, action, next_obs). Future-violation labels stop at episode boundaries. Set cost_shaping_enabled=False for the old raw-cost ablation.
  • The actor's cost advantages are mean-centred by default to remove the sampled baseline error. Optional full normalisation also divides the constraint residual by the same standard deviation, preserving the geometry of the CPO subproblem instead of silently changing the relative scale of b and c.
  • Reward and cost GAE use the value of every transition's actual next observation. A truncation bootstraps that next value but stops the GAE recurrence at the episode boundary; a true termination does neither. This avoids both the old rollout-tail approximation and leakage between episodes contained in one rollout.

Interface:

  • ICPO(cfg, seed=0) retains the learned C+ shaping extension.
  • CPO(cfg=None, seed=0) is the plain, standard CPO preset. It disables learned cost shaping and is exported alongside CPO_DEFAULT_CONFIG from src.rl.
  • Training uses the same modern contract as IPPO and IPPO_Lagrangian: trainer.train(env, timesteps) -> TrainResult.
  • The primary experiment notebooks now instantiate the plain CPO preset and publish CPO, run_cpo, and cpo_histories.pkl. ICPO remains available for an explicit shaped-cost ablation, but it is no longer labelled as the paper's plain CPO baseline.

Constraint estimate note:

  • The local trainer uses completed episodic cost from true episode starts minus cost_limit as the constraint violation signal. The estimator persists its partial episode across rollout boundaries, so a rollout fragment is never mistaken for a fresh episode.
  • cost_budget_discount_factor controls this budget estimator independently of the cost critic's cost_discount_factor. Its practical default is 1.0, matching the undiscounted episodic cost limits used by common safe-RL benchmarks. Discounted constraint objectives remain available explicitly.
  • If a rollout contains no completed episode, the most recent complete-episode estimate is retained. Before the first completion, ICPO initializes the shaped estimate at cost_limit, giving a neutral residual rather than treating an incomplete fragment as safe or unsafe.
  • The analytic CPO subproblem and line search use a scaled copy of that budget violation so its units match the mean cost-advantage surrogate. The scale is 1 / mean(sum_t cost_budget_discount_factor ** t) over completed episodes. For long discounted episodes this approaches the paper's (1 - gamma) normalization; for undiscounted finite episodes it becomes 1 / mean_episode_length.
  • mean_cost now refers to the shaped CPO cost used by the constraint. Raw monitor costs remain visible through raw_mean_cost and raw_cost_return_budget.
  • Metrics expose both the raw constraint_violation and the scaled_constraint_violation actually passed to the CPO solve.
  • cost_budget_fresh and completed_episodes make it explicit whether an update received a new complete-episode constraint estimate.
  • Additional shaping diagnostics include shaped_mean_cost, shaped_cost_return_budget, cost_shaping_delta_mean, cost_shaping_label_mean, and cost_shaping_loss.