Multi-player multi-shaper runners and lola (#166)

* most of the changes * most of the changes * more stuff
ucl-dark · Oct 11, 2023 · 7936b16 · 7936b16
1 parent 4665b7b
commit 7936b16
Show file tree

Hide file tree

Showing 29 changed files with 6,340 additions and 123 deletions.
diff --git a/pax/agents/agent.py b/pax/agents/agent.py
@@ -1,8 +1,9 @@
 from typing import Tuple
 
-from pax.utils import MemoryState, TrainingState
 import jax.numpy as jnp
 
+from pax.utils import MemoryState, TrainingState
+
 
 class AgentInterface:
     """Interface for agents to interact with runners and environemnts.

diff --git a/pax/agents/hyper/ppo.py b/pax/agents/hyper/ppo.py
@@ -329,7 +329,9 @@ def model_update_epoch(
             return new_state, new_mem, metrics
 
         @jax.jit
-        def make_initial_state(key: Any, hidden: jnp.ndarray) -> TrainingState:
+        def make_initial_state(
+            key: Any, hidden: jnp.ndarray
+        ) -> Tuple[TrainingState, MemoryState]:
             """Initialises the training state (parameters and optimiser state)."""
             key, subkey = jax.random.split(key)
             dummy_obs = jnp.zeros(shape=obs_spec)

diff --git a/pax/agents/lola/__init__.py b/pax/agents/lola/__init__.py