Skip to content

Allow experiment-local joint environments in training and evaluation - #36

Draft
PaulHax wants to merge 1 commit into
cotraining-ci-cleanupfrom
experiment-env-contract
Draft

PaulHax wants to merge 1 commit into
cotraining-ci-cleanupfrom
experiment-env-contract

Conversation

@PaulHax

@PaulHax PaulHax commented Oct 26, 2026 •

Copy link
Copy Markdown
Collaborator

Internal environment construction prevents local rules from sharing the rollout path

An experiment-local rule needs to reach both training and evaluation through the
same construction boundary. Supplying a new Python rule is insufficient when the
shared loops construct their environments internally or invoke different
transition methods.

Accept an explicit builder and validate the supported joint contract

Add an optional environment construction callable to joint training and
learned-policy matchup evaluation. Construct outside compiled rollout and check
reset, step, batched paths, and action masks against the canonical joint CC4 JAX
structure. Existing callers retain default construction.

Include a reward-offset example that changes Blue's payoff, applies the opposite
change to Red, and keeps raw reward components aligned. It demonstrates a
construction mechanism without introducing a research result.

Support remains bounded by the canonical JAX zero-sum structure

Custom environments must preserve ordered teams, action meanings and dimensions,
observation dimensions, state and constant pytrees, dtypes, and masks. Transformed
behavior must use explicit state and PRNG keys. Shared team rewards must remain
zero sum. Structural checks cannot prove purity or scientific equivalence. New
encodings, layouts, architectures, backends, and automatic execution-record
capture require further work.

Related: #30.

Depends on #38, which repairs inherited co-training CI checks.

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.

1 participant