Add support for loading PufferLib policies (Metta-AI#233)

# Add support for loading PufferLib policies This PR adds support for loading PufferLib policies from local files using the `puffer://` URI scheme. It includes: - New `PolicyState` tensorclass to manage policy state - Support for loading PufferLib policies in the policy store - Implementation of PufferLib policy and recurrent wrapper classes - Updates to the simulator to work with PufferLib policies - Refactoring of the play tool to separate policy loading from simulation - Simplifying eval configs to be shorter by removing redundant fields - Removes e3b for now, we have a plan to add it in later, cleaner The changes allow for seamless integration of policies trained with PufferLib into the Metta environment.

Author: daveey

configs/replay_job.yaml

@@ -2,7 +2,6 @@ defaults:
  - common
  - wandb: metta_research
  - sim: simple
- - /sim/simple@replay_job.sim
  - _self_
 
 cmd: play
@@ -13,4 +12,5 @@ torch_deterministic: true
 eval_db_uri: null
 
 replay_job:
+ sim: ${sim}
  policy_uri: ${policy_uri}

configs/sim/all.yaml

@@ -2,93 +2,61 @@ defaults:
  - sim
  - _self_
 
-num_envs: 1
-num_episodes: 1
-run_dir: ${run_dir}
-
 simulations:
  navigation/emptyspace_withinsight:
  env: env/mettagrid/navigation/evals/emptyspace_withinsight
- policy_agents_pct: 1.0
  navigation/emptyspace_outofsight:
  env: env/mettagrid/navigation/evals/emptyspace_outofsight
- policy_agents_pct: 1.0
  navigation/emptyspace_sparse:
  env: env/mettagrid/navigation/evals/emptyspace_sparse
- policy_agents_pct: 1.0
  navigation/walls_withinsight:
  env: env/mettagrid/navigation/evals/walls_withinsight
- policy_agents_pct: 1.0
  navigation/walls_outofsight:
  env: env/mettagrid/navigation/evals/walls_outofsight
- policy_agents_pct: 1.0
  navigation/walls_sparse:
  env: env/mettagrid/navigation/evals/walls_sparse
- policy_agents_pct: 1.0
  navigation/cylinder:
  env: env/mettagrid/navigation/evals/cylinder
- policy_agents_pct: 1.0
  navigation/obstacles0:
  env: env/mettagrid/navigation/evals/obstacles0
- policy_agents_pct: 1.0
  navigation/obstacles1:
  env: env/mettagrid/navigation/evals/obstacles1
- policy_agents_pct: 1.0
  navigation/obstacles2:
  env: env/mettagrid/navigation/evals/obstacles2
- policy_agents_pct: 1.0
  navigation/obstacles3:
  env: env/mettagrid/navigation/evals/obstacles3
- policy_agents_pct: 1.0
  navigation/corridors:
  env: env/mettagrid/navigation/evals/corridors
- policy_agents_pct: 1.0
  navigation/labyrinth:
  env: env/mettagrid/navigation/evals/labyrinth
- policy_agents_pct: 1.0
  navigation/radialmaze:
  env: env/mettagrid/navigation/evals/radialmaze
- policy_agents_pct: 1.0
  object_use/altar_use_free:
  env: env/mettagrid/object_use/evals/altar_use_free
- policy_agents_pct: 1.0
  object_use/altar_use:
  env: env/mettagrid/object_use/evals/altar_use
- policy_agents_pct: 1.0
  object_use/armory_use_free:
  env: env/mettagrid/object_use/evals/armory_use_free
- policy_agents_pct: 1.0
  object_use/armory_use:
  env: env/mettagrid/object_use/evals/armory_use
- policy_agents_pct: 1.0
  object_use/generator_use_free:
  env: env/mettagrid/object_use/evals/generator_use_free
- policy_agents_pct: 1.0
- object_use/generator_uses:
+ object_use/generator_use:
  env: env/mettagrid/object_use/evals/generator_use
- policy_agents_pct: 1.0
  object_use/lasery_use_free:
  env: env/mettagrid/object_use/evals/lasery_use_free
- policy_agents_pct: 1.0
  object_use/lasery_use:
  env: env/mettagrid/object_use/evals/lasery_use
- policy_agents_pct: 1.0
  object_use/mine_use:
  env: env/mettagrid/object_use/evals/mine_use
- policy_agents_pct: 1.0
  object_use/shoot_out:
  env: env/mettagrid/object_use/evals/shoot_out
- policy_agents_pct: 1.0
  object_use/swap_in:
  env: env/mettagrid/object_use/evals/swap_in
- policy_agents_pct: 1.0
  object_use/swap_out:
  env: env/mettagrid/object_use/evals/swap_out
- policy_agents_pct: 1.0
  object_use/temple_use_free:
  env: env/mettagrid/object_use/evals/temple_use_free
- policy_agents_pct: 1.0
  simple_npc:
  env: env/mettagrid/simple
- policy_agents_pct: 0.5
  npc_policy_uri: wandb://run/b.daveey.t.8.rdr9.3

configs/sim/memory.yaml

@@ -2,17 +2,10 @@ defaults:
  - sim
  - _self_
 
-num_envs: 1
-num_episodes: 1
-run_dir: ${run_dir}
-
 simulations:
  navigation/memory_easy:
  env: env/mettagrid/memory/evals/memory_easy
- policy_agents_pct: 1.0
  navigation/memory_medium:
  env: env/mettagrid/memory/evals/memory_medium
- policy_agents_pct: 1.0
  navigation/memory_hard:
  env: env/mettagrid/memory/evals/memory_hard
- policy_agents_pct: 1.0

configs/sim/navigation.yaml

@@ -2,53 +2,35 @@ defaults:
  - sim
  - _self_
 
-num_envs: 1
-num_episodes: 1
-run_dir: ${run_dir}
-
 simulations:
  navigation/emptyspace_withinsight:
  env: env/mettagrid/navigation/evals/emptyspace_withinsight
- policy_agents_pct: 1.0
  navigation/emptyspace_outofsight:
  env: env/mettagrid/navigation/evals/emptyspace_outofsight
- policy_agents_pct: 1.0
  navigation/emptyspace_sparse:
  env: env/mettagrid/navigation/evals/emptyspace_sparse
- policy_agents_pct: 1.0
  navigation/walls_withinsight:
  env: env/mettagrid/navigation/evals/walls_withinsight
- policy_agents_pct: 1.0
  navigation/walls_outofsight:
  env: env/mettagrid/navigation/evals/walls_outofsight
- policy_agents_pct: 1.0
  navigation/walls_sparse:
  env: env/mettagrid/navigation/evals/walls_sparse
- policy_agents_pct: 1.0
  navigation/cylinder:
  env: env/mettagrid/navigation/evals/cylinder
- policy_agents_pct: 1.0
  navigation/obstacles0:
  env: env/mettagrid/navigation/evals/obstacles0
- policy_agents_pct: 1.0
  navigation/obstacles1:
  env: env/mettagrid/navigation/evals/obstacles1
- policy_agents_pct: 1.0
  navigation/obstacles2:
  env: env/mettagrid/navigation/evals/obstacles2
- policy_agents_pct: 1.0
  navigation/obstacles3:
  env: env/mettagrid/navigation/evals/obstacles3
- policy_agents_pct: 1.0
  navigation/corridors:
  env: env/mettagrid/navigation/evals/corridors
- policy_agents_pct: 1.0
  navigation/labyrinth:
  env: env/mettagrid/navigation/evals/labyrinth
- policy_agents_pct: 1.0
  navigation/radialmaze:
  env: env/mettagrid/navigation/evals/radialmaze
- policy_agents_pct: 1.0
  navigation/cylinder_easy:
  env: env/mettagrid/navigation/evals/cylinder_easy
  navigation/honeypot:

configs/sim/object_use.yaml

@@ -2,47 +2,30 @@ defaults:
  - sim
  - _self_
 
-num_envs: 10
-num_episodes: 10
-run_dir: ${run_dir}
-
 simulations:
  object_use/altar_use_free:
  env: env/mettagrid/object_use/evals/altar_use_free
- policy_agents_pct: 1.0
  object_use/altar_use:
  env: env/mettagrid/object_use/evals/altar_use
- policy_agents_pct: 1.0
  object_use/armory_use_free:
  env: env/mettagrid/object_use/evals/armory_use_free
- policy_agents_pct: 1.0
  object_use/armory_use:
  env: env/mettagrid/object_use/evals/armory_use
- policy_agents_pct: 1.0
  object_use/generator_use_free:
  env: env/mettagrid/object_use/evals/generator_use_free
- policy_agents_pct: 1.0
  object_use/generator_use:
  env: env/mettagrid/object_use/evals/generator_use
- policy_agents_pct: 1.0
  object_use/lasery_use_free:
  env: env/mettagrid/object_use/evals/lasery_use_free
- policy_agents_pct: 1.0
  object_use/lasery_use:
  env: env/mettagrid/object_use/evals/lasery_use
- policy_agents_pct: 1.0
  object_use/mine_use:
  env: env/mettagrid/object_use/evals/mine_use
- policy_agents_pct: 1.0
  object_use/shoot_out:
  env: env/mettagrid/object_use/evals/shoot_out
- policy_agents_pct: 1.0
  object_use/swap_in:
  env: env/mettagrid/object_use/evals/swap_in
- policy_agents_pct: 1.0
  object_use/swap_out:
  env: env/mettagrid/object_use/evals/swap_out
- policy_agents_pct: 1.0
  object_use/temple_use_free:
  env: env/mettagrid/object_use/evals/temple_use_free
- policy_agents_pct: 1.0

configs/sim/sim.yaml

@@ -2,3 +2,6 @@ device: ${device}
 vectorization: ${vectorization}
 eval_db_uri: ${eval_db_uri}
 env: env/mettagrid/simple
+run_dir: ${run_dir}
+num_envs: 1
+num_episodes: 1

configs/sim/simple.yaml

@@ -2,8 +2,6 @@ defaults:
  - sim
  - _self_
 
-# This is a single simulation, not a suite.
-env: env/mettagrid/simple
-policy_agents_pct: 1.0
-num_envs: 1
-num_episodes: 1
+simulations:
+ simple:
+ env: env/mettagrid/simple

configs/sim/smoke_test.yaml

@@ -4,11 +4,7 @@ defaults:
  - sim
  - _self_
 
-num_envs: 1
-num_episodes: 1
-run_dir: ${run_dir}
 
 simulations:
  emptyspace_withinsight:
  env: env/mettagrid/navigation/evals/emptyspace_withinsight
- policy_agents_pct: 1.0

configs/user/daveey.yaml

@@ -1,9 +1,9 @@
 # @package __global__
 
 defaults:
- - override /agent: simple
- - override /eval: simple_solo
- - override /analyzer: eval_analyzer
+ - override /sim: simple
+#  - override /eval: simple_solo
+#  - override /analyzer: eval_analyzer
  - _self_
 
 trainer:
@@ -23,35 +23,34 @@ trainer:
 # policy_uri: wandb://run/b.daveey.t.1.bl
 # policy_uri: wandb://run/b.daveey.t.16.dr0
 
-# policy_uri: wandb://run/b.daveey.t.4.lra.muon
-policy_uri: wandb://run/b.daveey.t.1.lra.dr.muon
-# policy_uri: /tmp/puffer_metta_serial.pt
+# policy_uri: wandb://run/b.daveey.dr9.muon.latest
+# policy_uri: wandb://run/b.daveey.t.1.lra.dr.muon
+policy_uri: puffer:///tmp/puffer_metta.pt
 
 npc_policy_uri: ${policy_uri}
 # npc_policy_uri: wandb://run/b.daveey.t.16.dr0
 # policy_uri: ${trained_policy_uri}
 # npc_policy_uri: ${trained_policy_uri}
 # eval_db_uri: wandb://artifacts/daveey_eval_testing
 
+eval_db_uri: ${run_dir}/eval_stats
+
 analyzer:
  policy_uri: ${..policy_uri}
- eval_stats_uri: ${run_dir}/eval_stats
+ view_type: latest
  analysis:
  metrics:
  - metric: episode_reward
  - metric: "heart.get"
 
-eval:
+
+sim:
  env: /env/mettagrid/puffer
- eval_db_uri: ${run_dir}/eval_stats
- num_envs: 10
- num_episodes: 10
+ num_envs: 1
+ num_episodes: 1
  max_time_s: 600
  # policy_agents_pct: 1
 
- policy_uri: ${..policy_uri}
- npc_policy_uri: ${..npc_policy_uri}
- # eval_db_uri: ${..eval_db_uri} #file://daphne/sweep_stats
  # env: /env/mettagrid/reward_dr
  # env_overrides:
  # # sampling: 0.7
@@ -66,7 +65,7 @@ eval:
 wandb:
  checkpoint_interval: 1
 
-run_id: 16
+run_id: 17
 run: ${oc.env:USER}.local.${run_id}
 trained_policy_uri: ${run_dir}/checkpoints
 

metta.code-workspace

@@ -197,6 +197,40 @@
 }
 },
 "terminal.integrated.cwd": "${workspaceFolder}",
-"terminal.integrated.splitCwd": "workspaceRoot"
+"terminal.integrated.splitCwd": "workspaceRoot",
+"cSpell.words": [
+"bptt",
+"clipfrac",
+"coef",
+"cooldown",
+"dones",
+"heavyball",
+"kickstarter",
+"Kruskal",
+"lasery",
+"lexsort",
+"logratio",
+"metta",
+"mettagrid",
+"minibatch",
+"minibatches",
+"newlogprob",
+"nvec",
+"pufferlib",
+"pytest",
+"PYTHONPATH",
+"raylib",
+"relu",
+"tensorclass",
+"tensordict",
+"timestep",
+"timesteps",
+"truncateds",
+"unclipped",
+"vecenv",
+"venv",
+"vloss",
+"wandb"
+]
 }
 }