Autorom updates

README.md

@@ -1,6 +1,6 @@
 # MIXTAPE
 
-## Develop with Docker (recommended quickstart)
+## Develop with Docker
 This is the simplest configuration for developers to start with.
 
 ### Initial Setup
@@ -25,34 +25,67 @@ To non-destructively update your development stack at any time:
 
 ## Add data
 
-### Training
-Start by training an environment of your choice. For example:
+### Training - Supported environments
+
+- **Knights Archers Zombies** (`knights_archers_zombies_v10`, PettingZoo)
+  - Multi‑agent, discrete actions. Demonstrates combat/strategy coordination.
+  - Example:
+    ```bash
+    docker compose run --rm django \
+    ./manage.py training -e knights_archers_zombies_v10 \  # Environment
+    -a PPO \                                               # Agent
+    -p \                                                   # Parallel
+    -g 0.0 \                                               # GPUs
+    -t 100 \                                               # Iterations
+    --immediate
+    ```
+
+- **Pistonball** (`pistonball_v6`, PettingZoo)
+  - Multi‑agent with a continuous action space. Demonstrates continuous control and teamwork.
+  - Example:
+    ```bash
+    docker compose run --rm django \
+    ./manage.py training -e pistonball_v6 \  # Environment
+    -a PPO \                                 # Agent
+    -g 0.0 \                                 # GPUs
+    -t 100 \                                 # Iterations
+    --immediate
+    ```
+
+- **LunarLander** (`LunarLander-v2`, Gymnasium)
+  - Single‑agent, discrete actions. Demonstrates balancing multiple variables to achieve a safe, stable landing. Ideal candidate for decomposed rewards.
+  - Example:
+    ```bash
+    docker compose run --rm django \
+    ./manage.py training -e LunarLander-v2 \  # Environment
+    -a PPO \                                  # Agent
+    -g 0.0 \                                  # GPUs
+    -t 100 \                                  # Iterations
+    --immediate
+    ```
+
+Notes:
+- Use `-p/--parallel` only for PettingZoo environments.
+- DQN is for discrete action spaces; it is not available for Pistonball (continuous).
+
+### Inference
+
+Review available checkpoints:
 ```bash
-docker compose run --rm django \
-./manage.py training \
--e knights_archers_zombies_v10 \
--a PPO \
--p \
--g 0.0 \
--t 100 \
---immediate
-
-# For a detailed breakdown of all available options, use -h|--help
-docker compose run --rm django ./manage.py training --help
+docker compose run --rm django ./manage.py list_checkpoints
 ```
 
-### Inference
+You will see a list of available checkpoints, with the most recent at the top.
+```bash
+environment                 | checkpoint_pk | created             | inferences | episodes
+----------------------------+---------------+---------------------+------------+---------
+pistonball_v6               | 2             | 2026-01-03 19:29:25 | 1          | 1
+knights_archers_zombies_v10 | 1             | 2026-01-03 19:27:32 | 1          | 1
+```
 
 Select an existing checkpoint to run inference. For example:
 ```bash
-docker compose run --rm django \
-./manage.py inference \
-1 \
--p \
---immediate
-
-# For a detailed breakdown of all available options, use -h|--help
-docker compose run --rm django ./manage.py inference --help
+docker compose run --rm django ./manage.py inference 2 -p --immediate
 ```
 
 If you've already started the server with `docker compose up`, you can see all available checkpoints at <http://localhost:8000/admin/core/checkpoint/>.
@@ -69,6 +102,9 @@ When running the "Develop with Docker" configuration, all tox commands must be r
 ### Running Tests
 Run `uv run tox` to launch the full test suite.
 
+Note: Tests are configured to run with `--reuse-db` by default. If you change migrations or suspect the test database is out of sync, rebuild with:
+`docker compose run --rm django uv run tox -e test -- --create-db`
+
 Individual test environments may be selectively run.
 This also allows additional options to be be added.
 Useful sub-commands include:

mixtape/core/analysis/ray_utils/environments.py

@@ -8,7 +8,14 @@
 import supersuit as ss
 
 from mixtape.core.analysis.ray_utils.wrappers import ParallelPZWrapper, PZWrapper
-from mixtape.core.models.training import ExampleEnvs
+
+
+def is_gymnasium_registered_env(env_name: str) -> bool:
+    try:
+        gym.spec(env_name)
+    except Exception:
+        return False
+    return True
 
 
 def reshape_if_necessary(env) -> AECEnv | ParallelEnv:
@@ -39,14 +46,14 @@ def parallel_env_creator(env_module: types.ModuleType, config: dict) -> Parallel
 
 def gym_env_creator(env_module: str, config: dict) -> gym.Env:
     # Create the selected Gymnasium environment
-    return gym.make(f'ale_py:ALE/{env_module}', render_mode='rgb_array', **config)
+    return gym.make(env_module, render_mode='rgb_array', **config)
 
 
 def register_environment(
     env_name: str, config: dict, parallel: bool
 ) -> AECEnv | ParallelEnv | gym.Env:
     # Register the selected environment with RLlib
-    if ExampleEnvs.is_gymnasium_env(env_name):
+    if is_gymnasium_registered_env(env_name):
         register_env(env_name, lambda config: gym_env_creator(env_name, config))
         return gym_env_creator(env_name, config)
     else:

mixtape/core/management/commands/list_checkpoints.py

@@ -0,0 +1,53 @@
+from django.db.models import Count
+from django.utils import timezone
+import djclick as click
+
+from mixtape.core.models.checkpoint import Checkpoint
+
+
+@click.command()
+def list_checkpoints() -> None:
+    """List available checkpoints."""
+    checkpoints = Checkpoint.objects.select_related('training')
+
+    checkpoints = checkpoints.annotate(
+        inference_count=Count('inference', distinct=True),
+        episode_count=Count('inference__episode', distinct=True),
+    ).order_by('-created')
+
+    checkpoints_list = list(checkpoints)
+
+    headers = ('environment', 'checkpoint_pk', 'created', 'inferences', 'episodes')
+    rows: list[tuple[str, str, str, str, str]] = []
+
+    for checkpoint in checkpoints_list:
+        created_local = timezone.localtime(checkpoint.created)
+        rows.append(
+            (
+                checkpoint.training.environment,
+                str(checkpoint.pk),
+                created_local.strftime('%Y-%m-%d %H:%M:%S'),
+                str(checkpoint.inference_count),
+                str(checkpoint.episode_count),
+            )
+        )
+
+    if not rows:
+        click.echo('No checkpoints found.')
+        return
+
+    col_widths = [len(header) for header in headers]
+    for row in rows:
+        for idx, value in enumerate(row):
+            col_widths[idx] = max(col_widths[idx], len(value))
+
+    def format_row(values: tuple[str, ...]) -> str:
+        return ' | '.join(value.ljust(col_widths[idx]) for idx, value in enumerate(values))
+
+    header_line = format_row(headers)
+    separator_line = '-+-'.join('-' * width for width in col_widths)
+
+    click.echo(header_line)
+    click.echo(separator_line)
+    for row in rows:
+        click.echo(format_row(row))

mixtape/core/models/episode.py

@@ -1,6 +1,5 @@
 from typing import Any
 
-from django.conf import settings
 from django.db import models, transaction
 from django.db.models.signals import post_save
 from django.dispatch import receiver
@@ -17,7 +16,7 @@ def auto_compute_clustering(
     sender: type[Episode], instance: Episode, created: bool, **_: Any
 ) -> None:
     """Compute clustering when new episodes are created."""
-    if created and not getattr(settings, 'TESTING', False):
+    if created:
         # Avoid circular import
         from mixtape.core.tasks.clustering_tasks import compute_single_episode_clustering
 

mixtape/core/ray_utils/constants.py

@@ -0,0 +1,26 @@
+from django.db import models
+
+
+class ExampleEnvs(models.TextChoices):
+    # Example PettingZoo Environments
+    PZ_KnightsArchersZombies = 'knights_archers_zombies_v10'
+    PZ_Pistonball = 'pistonball_v6'
+    # Example Gymnasium Environments
+    GYM_LunarLander = 'LunarLander-v2'
+
+    @classmethod
+    def type(cls, value) -> Literal['PettingZoo', 'Gymnasium', 'Unknown']:
+        if cls(value).name.startswith('PZ'):
+            return 'PettingZoo'
+        elif cls(value).name.startswith('GYM'):
+            return 'Gymnasium'
+        return 'Unknown'
+
+    @classmethod
+    def is_gymnasium_env(cls, env_name: str) -> bool:
+        return cls.type(env_name) == 'Gymnasium'
+
+
+class SupportedAlgorithm(models.TextChoices):
+    PPO = 'PPO'
+    DQN = 'DQN'

mixtape/core/tasks/inference_tasks.py

@@ -7,7 +7,6 @@
 
 from celery import shared_task
 from django.db import transaction
-import numpy as np
 
 from mixtape.core.models import AgentStep, Episode, Inference
 from mixtape.core.models.step import Step
@@ -20,8 +19,6 @@
 def run_inference_task(inference_pk: int):
     # Import slow / task-specific dependencies locally
     import gymnasium as gym
-    from gymnasium.wrappers.atari_preprocessing import AtariPreprocessing
-    from gymnasium.wrappers.frame_stack import FrameStack
     from pettingzoo import AECEnv
     from pettingzoo.utils import ParallelEnv
     from ray.rllib.algorithms.algorithm import Algorithm
@@ -30,10 +27,11 @@ def run_inference_task(inference_pk: int):
 
     inference = Inference.objects.select_related('checkpoint__training').get(pk=inference_pk)
     env_config = (inference.config or {}).get('env_config', {})
+    env_name = inference.checkpoint.training.environment
 
     with contextlib.closing(
         register_environment(
-            inference.checkpoint.training.environment,
+            env_name,
             env_config,
             inference.parallel,
         )
@@ -47,15 +45,12 @@ def run_inference_task(inference_pk: int):
                 episode = Episode.objects.create(inference=inference)
 
                 if isinstance(env, gym.Env):
-                    env = AtariPreprocessing(env, grayscale_obs=True, scale_obs=False, frame_skip=1)
-                    env = FrameStack(env, num_stack=4)
                     observation, _ = (
                         env.reset()
                     )  # `reset` returns observation for single agent in Gym envs
                     reward = 0.0
 
                     for step in itertools.count(start=0):
-                        observation = np.transpose(observation, (1, 2, 0))
                         action, state, extras = algorithm.compute_single_action(
                             observation, full_fetch=True
                         )
@@ -84,7 +79,8 @@ def run_inference_task(inference_pk: int):
                         agent_step.save()
 
                         # Step the environment forward with the computed actions
-                        observation, reward, terminated, truncated, info = env.step(action)
+                        observation, step_reward, terminated, truncated, info = env.step(action)
+                        reward = float(step_reward)
 
                         if terminated or truncated:
                             break

mixtape/core/tests/conftest.py

@@ -1,3 +1,6 @@
+from typing import Generator
+
+from click.testing import CliRunner
 from playwright.sync_api import BrowserContext
 import pytest
 from pytest_django.live_server_helper import LiveServer
@@ -13,9 +16,20 @@ def api_client() -> APIClient:
 # This intentionally overrides the built-in fixture from pytest_playwright.
 # This will also cause other built-in fixtures like "page" to have a base URL set.
 @pytest.fixture
-def context(live_server: LiveServer, new_context: CreateContextCallback) -> BrowserContext:
+def context(
+    live_server: LiveServer, new_context: CreateContextCallback
+) -> Generator[BrowserContext, None, None]:
     context = new_context(
         base_url=live_server.url,
     )
-    context.set_default_timeout(3_000)
-    return context
+    context.set_default_timeout(3000)
+    try:
+        yield context
+    finally:
+        context.close()
+
+
+@pytest.fixture
+def cli_runner() -> CliRunner:
+    # Don't catch exceptions, so they'll be raised in the test case
+    return CliRunner(catch_exceptions=False)

mixtape/core/tests/factories.py

@@ -1,5 +1,6 @@
 from collections import OrderedDict
 from io import BytesIO
+from pathlib import Path
 
 import PIL.Image
 import PIL.ImageDraw
@@ -37,6 +38,10 @@ class Meta:
 
     training = factory.SubFactory(TrainingFactory)
     last = True
+    archive = factory.django.FileField(
+        from_path=Path(__file__).parent / 'data' / 'kaz_checkpoint.tar.bz2',
+        filename='checkpoint/archive.tar.bz2',
+    )
 
 
 class InferenceFactory(factory.django.DjangoModelFactory[Inference]):