Checkpointing (#221)

Author: cswinter

configs/codecraft/arena_medium.ron

@@ -1,5 +1,6 @@
 // Achieves ~0.4 against eval opponent (old baseline: ~0.8-0.95): https://wandb.ai/entity-neural-network/enn-ppo/reports/Arena-Medium-baseline--VmlldzoxNzgwMTM1
 ExperimentConfig(
+    version: 1,
     env: (
         id: "CodeCraft",
         kwargs: "{\"objective\": \"ARENA_MEDIUM\", \"hardness\": 1.0, \"win_bonus\": 2.0, \"hidden_obs\": true}",

configs/xprun/train.ron

@@ -96,9 +96,6 @@ XpV0(
             env_secrets: {
                 "WANDB_API_KEY": "wandb-api-key",
             },
-            volumes: {
-                "/mnt/a/Dropbox/artifacts/xprun": "/mnt/xprun",
-            },
         )
     }
 )

configs/xprun/trainbc.ron

@@ -69,9 +69,6 @@ XpV0(
             env_secrets: {
                 "WANDB_API_KEY": "wandb-api-key",
             },
-            volumes: {
-                "/mnt/a/Dropbox/artifacts/xprun": "/mnt/xprun",
-            },
         )
     }
 )

configs/xprun/traincc.ron

@@ -96,9 +96,6 @@ XpV0(
             env_secrets: {
                 "WANDB_API_KEY": "wandb-api-key",
             },
-            volumes: {
-                "/mnt/a/Dropbox/artifacts/xprun": "/mnt/xprun",
-            },
         ),
 
         "codecraftserver": (

enn_ppo/enn_ppo/tests/test_training.py

@@ -1,9 +1,26 @@
+from hyperstate import StateManager
+
 from enn_ppo.config import RolloutConfig
-from enn_ppo.train import EnvConfig, OptimizerConfig, PPOConfig, TrainConfig, _train
+from enn_ppo.train import (
+    EnvConfig,
+    OptimizerConfig,
+    PPOConfig,
+    State,
+    TrainConfig,
+    initialize,
+    train,
+)
+from entity_gym.examples import ENV_REGISTRY
 from rogue_net.relpos_encoding import RelposEncodingConfig
 from rogue_net.rogue_net import RogueNetConfig
 
 
+def _train(cfg: TrainConfig) -> float:
+    sm = StateManager(TrainConfig, State, initialize, None)
+    sm._config = cfg
+    return train(sm, ENV_REGISTRY[cfg.env.id])
+
+
 def test_multi_armed_bandit() -> None:
     cfg = TrainConfig(
         total_timesteps=500,