Fix bug in MultiSnake environment (#215)

Author: cswinter

configs/entity-gym/multisnake_2snakes11len.ron

@@ -0,0 +1,33 @@
+// Achieves > 0.98 episodic return: https://wandb.ai/entity-neural-network/enn-ppo/reports/MultiSnake-2-snakes-11-length--VmlldzoxNzgwNzEw
+ExperimentConfig(
+    version: 0,
+    env: (
+        id: "MultiSnake",
+        kwargs: "{\"num_snakes\": 2, \"max_snake_length\": 11}",
+    ),
+    rollout: (
+        num_envs: 512,
+        steps: 128,
+        processes: 16,
+    ),
+    total_timesteps: 100000000,
+    net: (
+        d_model: 128,
+        n_layer: 2,
+        relpos_encoding: (
+            extent: [10, 10],
+            position_features: ["x", "y"],
+        ),
+    ),
+    optim: (
+        bs: 32768,
+        lr: 0.005,
+        max_grad_norm: 10,
+        micro_bs: 8192,
+    ),
+    ppo: (
+        ent_coef: 0.03,
+        gamma: 0.99,
+        anneal_entropy: true,
+    ),
+)

entity_gym/entity_gym/environment/environment.py

@@ -123,7 +123,7 @@ class Entity:
 
 @dataclass
 class ObsSpace:
-    entities: Dict[str, Entity]
+    entities: Dict[EntityType, Entity]
 
 
 @dataclass

entity_gym/entity_gym/examples/multi_snake.py

@@ -125,6 +125,7 @@ def act(self, action: Mapping[str, Action]) -> Observation:
         self.step += 1
         move_action = action["move"]
         self.last_scores = deepcopy(self.scores)
+        food_to_spawn = []
         assert isinstance(move_action, CategoricalAction)
         for id, move in move_action.items():
             snake = self.snakes[id]
@@ -153,7 +154,8 @@ def act(self, action: Mapping[str, Action]) -> Observation:
                             1.0 / (self.max_snake_length - 1) / self.num_snakes
                         )
                     self.food.pop(i)
-                    self._spawn_food(snake.color)
+                    # Don't spawn food immediately since it might spawn in front of another snake that hasn't moved yet
+                    food_to_spawn.append(snake.color)
                     break
             if not ate_food:
                 snake.segments = snake.segments[1:]
@@ -166,6 +168,8 @@ def act(self, action: Mapping[str, Action]) -> Observation:
                 ]
             ):
                 game_over = True
+        for color in food_to_spawn:
+            self._spawn_food(color)
         if self.step >= self.max_steps:
             game_over = True
         return self._observe(done=game_over)