Misc bugfixes (#237)

Author: cswinter

enn_ppo/enn_ppo/train.py

@@ -166,12 +166,14 @@ def train(
             xp_info.xp_def.project,
             xp_info.sanitized_name + "-" + xp_info.id,
         )
+        id = xp_info.id
         Path(str(out_dir)).mkdir(parents=True, exist_ok=True)
 
         init_process(xp_info)
         rank = xp_info.replica_index
         parallelism = xp_info.replicas()
     else:
+        id = None
         out_dir = None
         rank = 0
         parallelism = 1
@@ -253,6 +255,7 @@ def train(
                 name=run_name,
                 save_code=True,
                 dir=data_dir,
+                id=id,
             )
             wandb.watch(agent)
 

enn_zoo/enn_zoo/procgen_env/base_env.py

@@ -118,22 +118,21 @@ def observe(self) -> Observation:
                 [0],
             )
         }
+        total = 1
         for type_id, name in self._entity_types().items():
             feats = state.entities[state.entities[:, 6] == type_id]
             if feats.shape[0] > 0:
                 feats = np.concatenate(
                     [feats, global_feats.repeat(feats.shape[0], axis=0)],
                     axis=1,
                 )
+                total += feats.shape[0]
             else:
                 feats = np.zeros(
                     (0, feats.shape[1] + global_feats.shape[1]), dtype=np.float32
                 )
             entities[name] = feats
-        assert (
-            sum(e.features.shape[0] for e in entities.values())  # type: ignore
-            == state.entities.shape[0]
-        )
+        assert total == state.entities.shape[0]
 
         return Observation(
             entities=entities,

rogue_net/rogue_net/relpos_encoding.py

@@ -172,6 +172,7 @@ def __init__(
         if self.value_gate is not None:
             self.value_gate_proj = nn.Linear(dmodel, dhead)
         self.cached_rkvs: Optional[Tuple[torch.Tensor, torch.Tensor]] = None
+        self.global_entity = "__global__" in obs_space.entities
 
     def relattn_logits(self, queries: torch.Tensor) -> torch.Tensor:
         assert self.cached_rkvs is not None
@@ -213,6 +214,8 @@ def keys_values(
         # Type of each entity
         entity_type: torch.Tensor,
     ) -> Tuple[torch.Tensor, torch.Tensor]:
+        if "__global__" in x and not self.global_entity:
+            x = {k: v for k, v in x.items() if k != "__global__"}
         relative_positions, entity_type = self._relative_positions(
             x, index_map, packpad_index, shape, entity_type
         )