Skip to content

Implement Red/Blue Self-Play/Co-training for CEC Experiments - #23

Draft
Dmujt wants to merge 64 commits into
cecfrom
cotraining
Draft

Dmujt wants to merge 64 commits into
cecfrom
cotraining

Conversation

@Dmujt

@Dmujt Dmujt commented Aug 27, 2026 •

Copy link
Copy Markdown
Contributor

Summary

  • Core co-training and learned Red support
    • a35964b — Adds simultaneous Blue/Red training, frozen-opponent support, joint environments, multi-policy checkpoint bundles, and learned-policy matchups.
    • 1e36af6 — Restricts learned Red to observed information.
    • fcec351, c8ca0b5 — Fix agent initialization and align Red observation tests with discovered-host behavior.
    • 6bac6a8 — Fixes credit assignment across dormant Red periods and adds training diagnostics and historical-opponent evaluation.
  • Additional Agent Architectures, e.g., MAPPO
    • e38ccad — Adds RNN/LSTM policy support.
    • 93fc881 — Preserves recurrent sequences during PPO minibatching.
    • 6fd56ea — Adds per-team architecture overrides and propagates recurrent state through evaluation.
    • 563b6cc — Adds MAPPO training, centralized critics, evaluation integration, and tests.
    • 6e19db9 — Extends MAPPO support and experiment configurations to both teams.
  • Reproducible topology generation and sampling
    • 459c8f6 — Adds recipe-driven topology generation and separate training/evaluation seed ranges.
    • 776f71f — Connects pre-generated topology banks to evaluation.
    • c8a74df — Improves runtime topology generation performance.
    • 269ace7, 1269823 — Introduce sampling without replacement for training topology batches and update the documentation.
  • Add multi-step evaluation and more metrics
    • f16c573 — Adds evaluation against scripted FSM/rule-based Red agents.
    • 381ea36 — Supports multiple ordered post-training evaluation jobs and standardizes evaluation argument naming.
    • 2239196 — Adds confidentiality, integrity, and availability resilience metrics alongside reward.
    • 9bdddb3 — Adds cross-seed evaluation and organizes post-training evaluation scripts.
    • e1451f8 — Adds support for rerunning an individual evaluation.
    • 072514d — Adds native CybORG benchmarks to post-training workflows.
    • 5c03ddb — Adds fixed-FSM evaluation of checkpoints used in cross-play.
  • MLflow tracking and analysis
    • 7491534, d17daf8, 3159504 — Improves logging from remote runs.
    • abceb44, ef33422 — Fix MLflow port conflicts and duplicate-key errors.
    • b0df254, d20b62d — Add experiment-comparison plotting and combine CIA mean/std metrics in plots.
    • 6ddec51 — Adds co-training training/evaluation figures and analysis utilities.
  • Training and evaluation performance/correctness
    • be19833, 0c48a7c — Batch learned-policy and scripted-Red evaluation; reduce experiment evaluation workloads.
    • f163b21 — Enables GPU execution for post-training evaluation.
    • 7f7e3b5 — Reduces joint-environment reset and evaluation overhead, with corresponding recipe budget adjustments.
    • 318b97a — Masks unavailable Blue actions during scripted evaluation.
    • 68282ba — Adds a parity-check launcher
    • c7f823e, 6f6d7b9, 9bf99ae— Maintenance
  • Experiment recipes and comparisons
    • 2bdd4ff — Adds Blue/Red best-response experiment recipes.
    • 42d660a — Adds an experiment for asymmetric CC4 rule changes.
    • aa45515 — Updates IPPO/MAPPO environment-diversity comparison recipes.
    • a0ebe09 — Adds MAPPO experiments using pooled team observations instead of global-state critic inputs.
    • 9119576 — Consolidates co-training recipes under recipes/cotraining/ and updates references.
    • 4880c34 — Aligns recipe evaluation settings for consistent final runs.
    • 2b876e3 — Adds comparisons of standard and diversity-trained policies against the same learned Red opponents.
    • 4ab8019 — Matches recurrent and feedforward optimizer budgets.
  • SOTA Comparison
    • b5afc45 — Adds configurable Red reward and Blue traffic-blocking rules for investigating co-training collapse.
    • a0190b1 — Adds enhanced Blue evidence observations across training and evaluation for HMARL comparisons.

PaulHax and others added 11 commits May 15, 2026 15:21
JAX 0.4.38 silently no-ops the persistent compilation cache when any
jax.Array is materialized before the cache dir is set (jax-ml/jax#25768),
so every process paid full cold compile (~7 min GPU / ~20 min/cell CPU).
JAX 0.5.x has a worse bug: cache hits silently return stale .at[].set()
updates inside lax.fori_loop / lax.cond (jax-ml/jax#31733), corrupting
the env step — reproduced here with 17/35 blue-Remove tests failing on a
warm cache under 0.5.1. Bug is fixed from 0.7.1 onward; we land on 0.10.

Distrax is dropped because its transitive tensorflow-probability<=0.25
imports jax.interpreters.xla.pytype_aval_mappings, removed in JAX 0.7.
The 5 callsites only used distrax.Categorical for sampling / log_prob /
entropy, all easily expressed on jax.random.categorical and
jax.nn.log_softmax — see src/jaxborg/policies/categorical.py (31 lines,
flax struct so it's jit/vmap/scan-compatible). Dropping distrax also
removes TFP, gast, and decorator from the dep graph.

Knock-on bumps from JAX 0.7+ requiring ml-dtypes>=0.5 (numpy 2 C ABI):
  * numpy 1.26.4 -> 2.3.5 (overrides cyborg's pin; cyborg runtime is
    numpy-2-compatible empirically)
  * scipy 1.12 -> 1.17 (overrides jaxmarl's `scipy<=1.12` precautionary
    upper bound; cf. jaxmarl commit 3ffa5b8f and issue #175)
  * torch 2.2 -> 2.10 (overrides cyborg's pin; needed for numpy-2 init)

Also retires a flaky test that relied on accidental RNG alignment
between CybORG's numpy and JAX's threefry:
tests/subsystems/test_fsm_red_agent.py::test_fsm_hidden_state_applies_after_completion_step
walked both pipelines forward on seed=0 and asserted they hit the same
FSM sequence — a parity-by-coincidence that breaks any time JAX's PRNG
layout changes. Restructured as a pure-state check of the two-stage
delayed-update mechanism (fsm_red_schedule_post_step_update stages
red_fsm_delayed_states; fsm_red_apply_delayed_update commits it on the
next step), which is the actual invariant the test name promises.

Verification:
  * fast suite: 772/772, cold 154s -> warm 115s (cache hit, correct)
  * full suite incl. slow: 1214 passed, 100 skipped, 3 xfailed
    (xfails are pre-existing, unrelated)
  * cache miscompile probe: blue_remove.py 35/35 cold and 35/35 warm
    against the same cache dir — under 0.5.1 the same probe gave 17/35
    failures on warm
…erage

Remove three xfail tests in test_fsm_red_env_differential.py that depended
on retired CybORG green/replay tape infrastructure. Their first two checks
(red_4 known-hosts parity and red_4 action-selection parity) are already
covered by test_red_policy_matches_cyborg_multistep across 200 steps x 5
seeds. The third (end-state host_compromised/red_privilege parity under
FSM red + green phish) had no equivalent — existing green-sync tests use
SleepAgent for red, so no exploit/privesc chains fire.

Add TestFsmRedGreenSyncParity::test_no_critical_state_diffs_over_10_steps
which closes that gap via CC4DifferentialHarness(FSM red + EnterpriseGreen
+ sync_green_rng=True) and asserts at least one privesc fired so the test
can't pass on a degenerate trajectory. Also adds seed=0 to the existing
red_policy_parity parametrize to preserve the original tests' seed.
The 4h→6h→10h timeout bumps were chasing a non-existent ceiling: GitHub
hard-cancels any single job at 360 min regardless of timeout-minutes,
so timeout: 600 on main is a no-op. The slow gate needs to scale below
that line, not above it.

Switch test-slow to an 8-way matrix using pytest-split. 545 slow items
split 69×7 + 62; each shard runs well under an hour, wall-clock for
the gate drops from ~3.5h to ~1h once durations are populated.

.test_durations is held in the GH Actions cache (actions/cache@v4),
not git — committing it would thrash history every run since the
floats churn. Each shard restores the most recent cached durations,
runs with --store-durations, and saves its own updated copy under a
unique key. Over runs, fresh entries from every group accumulate.

Sized for the cold-start case where no durations exist yet:
- 8 shards (vs strictly-needed 4) buys headroom against round-robin
  clustering of heavy parametrized tests
- timeout-minutes: 240 — first-run safety net only; steady-state is
  well under an hour. Still under the 360 hard cap.

Adds pytest-split to the dev dep group; uv.lock regenerated.
The shard 8/8 hang was a JAX-array-as-int hot loop in
_global_host_to_relative_slot / _abs_subnet_to_relative — pure-Python
host-side helpers that were doing int(jax_array[i, j]) per element of
small const lookup tables. Each access round-tripped through full JAX
primitive dispatch (slice_p.bind → apply_primitive); on cold-cache CI
runners the JAX persistent-cache lock contention from the also-running
xdist workers turned that into an unbounded stall. Switched to one
np.asarray + np.flatnonzero per call. Module already declared itself
JAX-free at the top — this aligns with that intent.

Workflow:
- drop addopts via -o "addopts=" so -n auto doesn't stack on top of
  sharding (xdist + pytest-split contend on the JAX cache lock)
- fix the durations cache: previous strategy had each shard overwrite
  .test_durations with only its 1/N slice, so prefix-restore picked up
  partial files that never accumulated. Now each shard writes
  .test_durations_cache/group-N.json, and a merge step rebuilds the
  full .test_durations before pytest runs
- timeout 240 → 60 (hang root cause is fixed; no first-run cliff to
  budget for)

Verified locally: test_native_generic_exploit_respects_blocked_scan_-
source_route_matches_cyborg now finishes in 3m4s (was unbounded).
First run after the cache schema change has empty
.test_durations_cache, so pytest-split round-robins by test count and
heavy parametrized tests cluster into 1-2 shards that exceed 60 min
serially. Subsequent runs (with partials cached) all fit under 30 min.
Per-shard cache keys could never warm up: actions/cache restores at
most one cache per step, so each shard merged from a single shard's
2.4 KB partial instead of all eight. pytest-split fell back to
round-robin and shard 2/8 hit the 120 min timeout while shard 3/8
finished in 59 s.

Replace per-shard caches with a single aggregate refreshed by a new
aggregate-durations job that merges every shard's uploaded
.test_durations (plus the prior aggregate so timed-out shards keep
their last-known durations) and saves under
slow-durations-aggregated-<run_id>. Each shard restores the most
recent aggregate via the slow-durations-aggregated- prefix.

Bump test-slow timeout 120 -> 180 min for cold-cache headroom on the
first run before the aggregate exists.
Two fixes for the slow-test sharding bootstrap:

1. upload-artifact@v4 excludes hidden files by default, so
   `.test_durations` was silently dropped. The aggregate-durations
   job consequently found 0 artifacts and never populated the cache.
   Add `include-hidden-files: true`.

2. Even with #1 fixed, the very first run on a new cache scope
   (this PR, main after merge, future PRs) still cold-starts with
   round-robin sharding and hits the per-shard timeout. Add a
   conditional "seed from URL" step that pulls a known-good
   `.test_durations` from `vars.SLOW_DURATIONS_SEED_URL` when the
   cache is empty. Subsequent runs hit the cache and skip the
   download, and the aggregate job keeps the cache refreshed.
@Dmujt
Dmujt changed the base branch from main to cec August 27, 2026 03:47
Coordinate initial and automatic resets across parallel JAX environments. Require diversified banks to cover NUM_ENVS, preserve singleton controls, and support both Blue-only and joint trainers.
Extend Blue inputs from 210 to 402 with host presence, completed Analyse evidence, and persistent alert history. Match JAX and native CybORG memory semantics, propagate dimensions through policies and checkpoints, and enable the contract for dual-team cotraining recipes.
Evaluate final Blue policies against stock FSM Red for 100 native episodes. Allow each post-training job to select CPU or CUDA, and record simulator provenance and score semantics alongside the existing held-out JAX evaluations.
Use 16 whole-sequence minibatches for GRU and LSTM cotraining, yielding 64 optimizer steps per rollout. Remove unused CleanRL blocks and document that earlier eight-minibatch runs require rerunning for matched comparisons.
… topology; also updated invalid scripted-red settings and MAPPO checkpoint lookup
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants