Merge branch 'release' into rel-hpc

evaluate.py

@@ -62,7 +62,7 @@ def save_replays(policy_store_dir, save_dir):
     from reinforcement_learning import policy  # import your policy
     def make_policy(envs):
         learner_policy = policy.Baseline(
-            envs,
+            envs._driver_env,
             input_size=args.input_size,
             hidden_size=args.hidden_size,
             task_size=args.task_size

reinforcement_learning/policy.py

@@ -35,9 +35,11 @@ def critic(self, hidden):
 
 
 class Baseline(pufferlib.models.Policy):
-  def __init__(self, envs, input_size=256, hidden_size=256, task_size=4096):
+  def __init__(self, env, input_size=256, hidden_size=256, task_size=4096):
     super().__init__()
-    self.envs = envs
+
+    self.flat_observation_space = env.flat_observation_space
+    self.flat_observation_structure = env.flat_observation_structure
 
     self.tile_encoder = TileEncoder(input_size)
     self.player_encoder = PlayerEncoder(input_size, hidden_size)
@@ -50,7 +52,8 @@ def __init__(self, envs, input_size=256, hidden_size=256, task_size=4096):
     self.value_head = torch.nn.Linear(hidden_size, 1)
 
   def encode_observations(self, flat_observations):
-    env_outputs = self.envs.unpack_batched_obs(flat_observations)
+    env_outputs = pufferlib.emulation.unpack_batched_obs(flat_observations,
+        self.flat_observation_space, self.flat_observation_structure)
     tile = self.tile_encoder(env_outputs["Tile"])
     player_embeddings, my_agent = self.player_encoder(
         env_outputs["Entity"], env_outputs["AgentId"][:, 0]

train.py

@@ -31,7 +31,7 @@ def setup_env(args):
 
     def make_policy(envs):
         learner_policy = policy.Baseline(
-            envs,
+            envs._driver_env,
             input_size=args.input_size,
             hidden_size=args.hidden_size,
             task_size=args.task_size