init

no more .mpt

Merge remote-tracking branch 'origin/main' into richard-unifympt

slim

policy spec handler

more concise

cleanup

simplify

Merge remote-tracking branch 'origin/main' into richard-unifympt

re-add

fix policy spex

Update packages/mettagrid/python/src/mettagrid/util/uri_resolvers/schemes.py

Co-authored-by: graphite-app[bot] <96075541+graphite-app[bot]@users.noreply.github.com>

Merge remote-tracking branch 'origin/main' into richard-unifympt

bundles

Merge remote-tracking branch 'origin/main' into richard-unifympt

Merge remote-tracking branch 'origin/main' into richard-unifympt

bundle

Merge remote-tracking branch 'origin/richard-unifympt' into richard-unifympt

Merge remote-tracking branch 'origin/main' into richard-unifympt

Merge remote-tracking branch 'origin/main' into richard-unifympt

Merge remote-tracking branch 'origin/main' into richard-unifympt

Merge remote-tracking branch 'origin/main' into richard-unifympt

Merge remote-tracking branch 'origin/main' into richard-unifympt

Merge remote-tracking branch 'origin/main' into richard-unifympt

simplify?

ugh compat

cleanup

Merge remote-tracking branch 'origin/main' into richard-unifympt

cleanup

tests

Merge remote-tracking branch 'origin/main' into richard-unifympt

more tests

Merge remote-tracking branch 'origin/main' into richard-unifympt

Merge remote-tracking branch 'origin/main' into richard-unifympt

simplify?

Merge branch 'main' into richard-unifympt

cleanup

Merge remote-tracking branch 'origin/main' into richard-unifympt

Merge remote-tracking branch 'origin/main' into richard-unifympt

Merge remote-tracking branch 'origin/main' into richard-unifympt

no more .mpt

remove all .mpt and lint

cleanup

local data path fixes

mpt re-add

re-add artifact

lint

Merge remote-tracking branch 'origin/main' into richard-unifympt

more cleanup

Merge remote-tracking branch 'origin/main' into richard-unifympt

diff cleanup

ftt

lint

fix error

Merge remote-tracking branch 'origin/main' into richard-unifympt

more tests

lint

Merge remote-tracking branch 'origin/main' into richard-unifympt

Merge remote-tracking branch 'origin/main' into richard-unifympt

checkpoint policy does save/load

lint

checkpoint moving

catcus

lint

Merge branch 'main' into richard-unifympt

fold-in

[pyright 4] Get pyright to pass on app_backend (#4478)

Merge remote-tracking branch 'origin/main' into richard-unifympt

Fix command, add space (#4456)

added space to --app:lib--tlsEmulation:off which makes it --app:lib
--tlsEmulation:off
now it runs

Rename HyperUpdateRule to ScheduleRule (#4483)

## Summary
- rename HyperUpdateRule to ScheduleRule and apply to TrainerConfig via
target_path
- update recipes and teacher scheduling to use ScheduleRule
- report PPO stats using ppo_actor/ppo_critic hyperparam keys and update
tests

## Testing
- not run (not requested)

---------

Co-authored-by: graphite-app[bot] <96075541+graphite-app[bot]@users.noreply.github.com>

Merge remote-tracking branch 'origin/main' into richard-unifympt

Fix supervisor teacher behavior and legacy BC mode (#4484)

## Summary
- gate PPO actor during supervisor teacher phase
- fix supervisor/no-teacher behavior and add legacy BC (no gating, no
PPO resume)
- require supervisor policy URI for sliced_cloner_no_ppo

## Testing
- not run (not requested)

---------

Co-authored-by: graphite-app[bot] <96075541+graphite-app[bot]@users.noreply.github.com>
Co-authored-by: Adam S <134907338+gustofied@users.noreply.github.com>

Minor fixes to the slstm triton kernel, causing failures for certain kernel sizes (#4492)

cleanup

Merge remote-tracking branch 'origin/main' into richard-unifympt

fold in

training environments and eval environments mismatched (#4487)

I ran a direct config comparison using the training entrypoint
(recipes/experiment/cogs_v_clips.train) with variants=["heart_chorus"]
and compared the eval suite
config it builds (difficulty standard + heart_chorus) for an overlapping
mission: hello_world.oxygen_bottleneck.

  Findings:

  - Compass is ON in both training and eval (global_obs.compass=True).
- Vibe count and change‑vibe settings match (152 vibes;
change_vibe.number_of_vibes=152).
- But the mission parameters differ between training and eval for the
same mission name:
      - game.objects.carbon_extractor.max_uses: train 25 vs eval 100
      - game.objects.oxygen_extractor.max_uses: train 5 vs eval 20
      - game.objects.germanium_extractor.max_uses: train 5 vs eval 20
      - game.objects.silicon_extractor.max_uses: train 120 vs eval 480

So the mismatch isn’t compass — it’s the mission definitions used by
training vs eval. Training uses base missions
(cogames.cogs_vs_clips.missions), while eval uses
integrated eval missions (cogames.cogs_vs_clips.evals.integrated_evals)
that have different extractor settings.

Also: the eval suite used by recipes/experiment/cogs_v_clips.train does
not include machina_1.open_world at all (it only evaluates the 7
HELLO_WORLD integrated
evals). So training can be creating hearts on easier missions while your
eval runs on machina_1 are a different environment entirely.

  Relevant files:

- Compass default: packages/cogames/src/cogames/cogs_vs_clips/mission.py
- Training entrypoint + eval suite wiring:
recipes/experiment/cogs_v_clips.py
- Eval mission definitions:
packages/cogames/src/cogames/cogs_vs_clips/evals/integrated_evals.py

If you want true parity, we should align which mission templates eval
uses (and/or include machina_1.open_world in the eval suite). I can
patch this if you want —
  tell me whether you prefer:

1. Eval suite uses the same mission templates as training (from
missions.py), or
  2. Training uses the integrated eval mission definitions, or
  3. Add machina_1.open_world to the eval suite.

ripping out

Merge remote-tracking branch 'origin/main' into richard-unifympt

simplify

fix and lint

choke

simplify submission zip creation

use policy_spec for submission zips

tighten checkpoint io helpers

shorten checkpoint arg help

inline checkpoint policy helpers

restore policy spec docstring

validate checkpoint data_path before download

require checkpoint directory URIs

expand policy spec s3 docstring

Author: relh

metta/rl/loss/sl_checkpointed_kickstarter.py

@@ -159,6 +159,8 @@ def _construct_checkpoint_uri(self, epoch: int) -> str:
         filename = checkpoint_filename(run_name, epoch)
 
         if parsed.scheme == "file" and parsed.local_path:
+            if parsed.local_path.is_file():
+                raise ValueError("Provide a checkpoint directory, not policy_spec.json")
             path = parsed.local_path.parent / filename
             return f"file://{path}"
         elif parsed.scheme == "s3" and parsed.bucket and parsed.key:

metta/rl/metta_scheme_resolver.py

@@ -7,7 +7,6 @@
 from metta.app_backend.metta_repo import PolicyVersionWithName
 from metta.common.util.constants import PROD_STATS_SERVER_URI
 from mettagrid.util.uri_resolvers.base import MettaParsedScheme, SchemeResolver
-from mettagrid.util.uri_resolvers.schemes import resolve_uri
 
 logger = logging.getLogger(__name__)
 
@@ -118,12 +117,6 @@ def get_path_to_policy_spec_or_mpt(self, uri: str) -> str:
             logger.info(f"Metta scheme resolver: {uri} resolved to s3 policy spec: {policy_version.s3_path}")
             return policy_version.s3_path
 
-        # If that is missing (probably legacy policy), we send you to the mpt file, and will later assume
-        # that the class to hydrate from is MptPolicy
-        mpt_file_path = (policy_version.policy_spec or {}).get("init_kwargs", {}).get("checkpoint_uri")
-        if not mpt_file_path:
-            raise ValueError(f"Data not found for policy version {policy_version.id}")
-        if not mpt_file_path.endswith(".mpt"):
-            raise ValueError(f"Invalid mpt file path: {mpt_file_path}")
-        logger.info(f"Metta scheme resolver: {uri} resolved to mpt checkpoint: {mpt_file_path}")
-        return resolve_uri(mpt_file_path).canonical
+        raise ValueError(
+            f"Policy version {policy_version.id} has no s3_path; expected a policy spec submission zip in S3."
+        )

metta/rl/training/evaluator.py

@@ -7,6 +7,7 @@
 import os
 import uuid
 import zipfile
+from pathlib import Path
 from typing import Any, Optional
 
 import torch
@@ -27,9 +28,9 @@
 from metta.tools.utils.auto_config import auto_replay_dir
 from mettagrid.base_config import Config
 from mettagrid.policy.policy import PolicySpec
-from mettagrid.policy.submission import POLICY_SPEC_FILENAME
+from mettagrid.policy.submission import POLICY_SPEC_FILENAME, SubmissionPolicySpec
 from mettagrid.util.file import write_data
-from mettagrid.util.uri_resolvers.schemes import policy_spec_from_uri
+from mettagrid.util.uri_resolvers.schemes import policy_spec_from_uri, resolve_uri
 
 logger = logging.getLogger(__name__)
 
@@ -140,21 +141,34 @@ def should_evaluate(self, epoch: int) -> bool:
         return epoch % interval == 0
 
     def _create_submission_zip(self, policy_spec: PolicySpec) -> bytes:
-        """Create a submission zip containing policy_spec.json."""
+        """Create a submission zip containing policy_spec.json and optional weights."""
+        submission_spec = SubmissionPolicySpec.model_validate(policy_spec.model_dump(mode="json"))
+        data_path = submission_spec.data_path
+        if data_path and Path(data_path).is_absolute():
+            data_path = Path(data_path).name
+            submission_spec.data_path = data_path
+        spec_bytes = submission_spec.model_dump_json().encode("utf-8")
+
         buffer = io.BytesIO()
         with zipfile.ZipFile(buffer, "w", zipfile.ZIP_DEFLATED) as zipf:
-            zipf.writestr(POLICY_SPEC_FILENAME, policy_spec.model_dump_json())
+            zipf.writestr(POLICY_SPEC_FILENAME, spec_bytes)
+            if data_path:
+                if not policy_spec.data_path:
+                    raise ValueError("policy_spec.data_path missing for submission")
+                source_path = Path(policy_spec.data_path)
+                if not source_path.is_absolute():
+                    raise ValueError("policy_spec.data_path must be absolute for submission")
+                zipf.writestr(data_path, source_path.read_bytes())
         return buffer.getvalue()
 
-    def _upload_submission_zip(self, policy_spec: PolicySpec) -> str | None:
+    def _upload_submission_zip(self, policy_spec: PolicySpec, policy_uri: str) -> str | None:
         """Upload a submission zip to S3 and return the s3_path."""
-        checkpoint_uri = policy_spec.init_kwargs.get("checkpoint_uri")
-        if not checkpoint_uri or not checkpoint_uri.startswith("s3://"):
+        if not policy_uri.startswith("s3://"):
             return None
 
-        submission_path = checkpoint_uri.replace(".mpt", "-submission.zip")
-        zip_data = self._create_submission_zip(policy_spec)
-        write_data(submission_path, zip_data, content_type="application/zip")
+        checkpoint_dir = resolve_uri(policy_uri).canonical.rstrip("/")
+        submission_path = f"{checkpoint_dir}/submission.zip"
+        write_data(submission_path, self._create_submission_zip(policy_spec), content_type="application/zip")
         logger.info("Uploaded submission zip to %s", submission_path)
         return submission_path
 
@@ -163,6 +177,7 @@ def _create_policy_version(
         *,
         stats_client: StatsClient,
         policy_spec: PolicySpec,
+        policy_uri: str,
         epoch: int,
         agent_step: int,
     ) -> uuid.UUID:
@@ -176,7 +191,7 @@ def _create_policy_version(
         )
 
         # Upload submission zip to S3
-        s3_path = self._upload_submission_zip(policy_spec)
+        s3_path = self._upload_submission_zip(policy_spec, policy_uri)
 
         # Create policy version
         policy_version_id = stats_client.create_policy_version(
@@ -209,6 +224,7 @@ def evaluate(
             policy_version_id = self._create_policy_version(
                 stats_client=self._stats_client,
                 policy_spec=policy_spec,
+                policy_uri=policy_uri,
                 epoch=epoch,
                 agent_step=agent_step,
             )

metta/setup/shell.py

@@ -21,9 +21,8 @@ def help_configs() -> None:
     success("# Load configs with overrides:")
     info('cfg = load_cfg("train_job.yaml", ["training_env.curriculum=/env/mettagrid/arena/advanced"])')
     success("# Load checkpoints:")
-    info('artifact = load_mpt("file://./train_dir/my_run/checkpoints/my_run:v12.mpt")')
-    info('artifact = load_mpt("s3://bucket/path/my_run/checkpoints/my_run:v12.mpt")')
-    info('policy = artifact.instantiate(policy_env_info, torch.device("cpu"))')
+    info('spec = policy_spec_from_uri("file://./train_dir/my_run/checkpoints/my_run:v12")')
+    info("policy = initialize_or_load_policy(policy_env_info, spec)")
     success("# Create checkpoint manager:")
     info('cm = CheckpointManager(run="my_run", run_dir="./train_dir")')
 

packages/cogames/scripts/run_evaluation.py

@@ -14,9 +14,9 @@
   uv run python packages/cogames/scripts/run_evaluation.py \
       --agent cogames.policy.nim_agents.agents.ThinkyAgentsMultiPolicy --cogs 1
   uv run python packages/cogames/scripts/run_evaluation.py \
-      --agent cogames.policy.lstm.LSTMPolicy --checkpoint s3://bucket/path/model.mpt --cogs 1
+      --agent cogames.policy.lstm.LSTMPolicy --checkpoint s3://bucket/path/run:v<N> --cogs 1
   uv run python packages/cogames/scripts/run_evaluation.py \
-      --agent s3://bucket/path/model.mpt --cogs 1
+      --agent s3://bucket/path/run:v<N> --cogs 1
 """
 
 import argparse
@@ -38,6 +38,7 @@
 import matplotlib.pyplot as plt
 import numpy as np
 import torch
+from safetensors.torch import load as load_safetensors
 
 from cogames.cogs_vs_clips.evals.diagnostic_evals import DIAGNOSTIC_EVALS
 from cogames.cogs_vs_clips.mission import Mission, MissionVariant, NumCogsVariant
@@ -87,12 +88,13 @@ def _get_policy_action_space(policy_path: str) -> Optional[int]:
         return None
 
     try:
-        from mettagrid.policy.mpt_artifact import load_mpt
-
-        artifact = load_mpt(policy_path)
+        spec = policy_spec_from_uri(policy_path)
+        if not spec.data_path:
+            return None
+        weights = load_safetensors(Path(spec.data_path).read_bytes())
 
         # Look for actor head weight to determine action space
-        for key, tensor in artifact.state_dict.items():
+        for key, tensor in weights.items():
             if "actor_head" in key and "weight" in key and len(tensor.shape) == 2:
                 action_space = tensor.shape[0]
                 _policy_action_space_cache[policy_path] = action_space
@@ -1152,7 +1154,7 @@ def lookup_wrapper(_s: str, exp_name: str):
 def main():
     parser = argparse.ArgumentParser(description="Evaluate scripted or custom agents.")
     parser.add_argument("--agent", nargs="*", default=None, help="Agent key, class path, or S3 URI")
-    parser.add_argument("--checkpoint", type=str, default=None, help="Checkpoint path (or S3 URI)")
+    parser.add_argument("--checkpoint", help="Checkpoint directory URI")
     parser.add_argument("--experiments", nargs="*", default=None, help="Experiments to run")
     parser.add_argument("--variants", nargs="*", default=None, help="Variants to apply")
     parser.add_argument("--cogs", nargs="*", type=int, default=None, help="Agent counts to test")

packages/cogames/src/cogames/cli/policy.py

@@ -50,10 +50,10 @@ def list_checkpoints():
 
 def describe_policy_arg(with_proportion: bool):
     console.print("[bold cyan]-p [POLICY][/bold cyan] accepts two formats:\n")
-    console.print("[bold]1. URI format[/bold] (for .mpt checkpoints):")
+    console.print("[bold]1. URI format[/bold] (checkpoint bundle):")
     console.print("  - metta://policy/<name> or metta://policy/<uuid>")
-    console.print("  - s3://bucket/path/to/checkpoint.mpt")
-    console.print("  - file:///path/to/checkpoint.mpt or /path/to/checkpoint.mpt")
+    console.print("  - s3://bucket/path/to/run:v<N>")
+    console.print("  - file:///path/to/run:v<N> or /path/to/run:v<N>")
     console.print()
     console.print(
         "[bold]2. Key-value format[/bold]: "

packages/mettagrid/python/src/mettagrid/util/uri_resolvers/README.md

@@ -2,28 +2,30 @@
 
 This package provides a pluggable URI resolution system for handling different resource schemes.
 
+Checkpoint URIs point at a checkpoint directory containing `policy_spec.json`.
+
 ## Usage
 
 ```python
 from mettagrid.util.uri_resolvers.schemes import parse_uri, resolve_uri, get_checkpoint_metadata
 
 # Parse a URI to get its components
-parsed = parse_uri("s3://bucket/path/to/file.mpt")
+parsed = parse_uri("s3://bucket/path/to/run:v5")
 print(parsed.scheme)  # "s3"
 print(parsed.bucket)  # "bucket"
-print(parsed.key)     # "path/to/file.mpt"
+print(parsed.key)     # "path/to/run:v5"
 
 # Get checkpoint info (run_name, epoch) from parsed URI
 info = parsed.checkpoint_info  # ("run_name", 5) or None
 if info:
     run_name, epoch = info
 
 # Resolve a URI (normalizes and finds latest checkpoint if applicable)
-parsed = resolve_uri("file:///path/to/checkpoints")
-print(parsed.canonical)  # "file:///path/to/checkpoints/run:v5.mpt"
+parsed = resolve_uri("file:///path/to/checkpoints:latest")
+print(parsed.canonical)  # "file:///path/to/checkpoints/run:v5"
 
 # Get full checkpoint metadata (resolves URI first)
-metadata = get_checkpoint_metadata("s3://bucket/checkpoints/my-run:v5.mpt")
+metadata = get_checkpoint_metadata("s3://bucket/checkpoints/my-run:v5")
 print(metadata.run_name)  # "my-run"
 print(metadata.epoch)     # 5
 print(metadata.uri)       # resolved URI