Feature Implementation for Pressure Plates

.github/workflows/test-overcooked-v3.yml

@@ -0,0 +1,23 @@
+name: Overcooked V3 Tests
+
+on: [push, pull_request]
+
+jobs:
+  test:
+    runs-on: ubuntu-latest
+
+    steps:
+    - uses: actions/checkout@v4
+
+    - name: Set up Python 3.12
+      uses: actions/setup-python@v5
+      with:
+        python-version: "3.12"
+
+    - name: Install dependencies
+      run: |
+        python -m pip install --upgrade pip
+        pip install -e .[dev]
+
+    - name: Run overcooked_v3 tests
+      run: pytest tests/overcooked_v3/ -v

exports/custom_layout.txt

@@ -1,11 +1,13 @@
 # Add to jaxmarl/environments/overcooked_v3/layouts.py
 
 custom_layout = """
-A 01WWW
-A]]]]}2
-W{WWW}W
-W{[[[[W
-WBWWPWW
+WWWWWW
+0    W
+WWWW P
+1    X
+WWWW W
+2    W
+WWWWWW
 """
 
 # Add to overcooked_v3_layouts dictionary:

jaxmarl/environments/overcooked_v3/__init__.py

@@ -1,4 +1,5 @@
 """Overcooked V3 environment with pot burning, order queue, and conveyor belts."""
 
+from .common import ButtonAction
 from .layouts import overcooked_v3_layouts, Layout
 from .overcooked import OvercookedV3, ObservationType

jaxmarl/environments/overcooked_v3/common.py

@@ -29,6 +29,12 @@ class StaticObject(IntEnum):
     ITEM_CONVEYOR = 20
     PLAYER_CONVEYOR = 21
 
+    # Moving walls, buttons, barriers, pressure plates
+    MOVING_WALL = 22
+    BUTTON = 23
+    BARRIER = 24
+    PRESSURE_PLATE = 25
+
     @staticmethod
     def is_ingredient_pile(obj):
         return (obj >= StaticObject.INGREDIENT_PILE_BASE) & (obj < StaticObject.ITEM_CONVEYOR)
@@ -43,6 +49,18 @@ def ingredient_pile(idx):
         return StaticObject.INGREDIENT_PILE_BASE + idx
 
 
+class ButtonAction(IntEnum):
+    """Actions a button can trigger on its linked moving wall or barrier."""
+    TOGGLE_PAUSE = 0       # Pause/unpause the wall's movement
+    TOGGLE_DIRECTION = 1   # Reverse the wall's direction
+    TOGGLE_BOUNCE = 2      # Toggle bounce mode on/off
+    TRIGGER_MOVE = 3       # Move the wall one step (wall is paused by default)
+
+    # Barrier actions (button_linked_wall refers to barrier index instead)
+    TOGGLE_BARRIER = 4     # Toggle barrier active state
+    TIMED_BARRIER = 5      # Deactivate barrier temporarily (auto-reactivates)
+
+
 class DynamicObject(IntEnum):
     """Dynamic objects (channel 1) - bitwise encoding."""
     EMPTY = 0

jaxmarl/environments/overcooked_v3/interactive.py

@@ -0,0 +1,133 @@
+import argparse
+import jax
+import jax.numpy as jnp
+from jaxmarl.environments.overcooked_v3.common import Actions
+from jaxmarl.environments.overcooked_v3.overcooked import OvercookedV3
+from jaxmarl.environments.overcooked_v3.layouts import overcooked_v3_layouts as layouts
+from jaxmarl.viz.overcooked_v3_visualizer import OvercookedV3Visualizer
+
+
+class InteractiveOvercookedV3:
+
+    def __init__(self, layout, agent_view_size=None, no_jit=False, debug=False):
+        self.debug = debug
+        self.no_jit = no_jit
+
+        self.env = OvercookedV3(layout=layout, agent_view_size=agent_view_size)
+        self.viz = OvercookedV3Visualizer(self.env)
+
+    def run(self, key):
+        self.key = key
+        with jax.disable_jit(self.no_jit):
+            self._run()
+
+    def _run(self):
+        self._reset()
+
+        self.viz.window.reg_key_handler(self._handle_input)
+        self.viz.show(block=True)
+
+    def _handle_input(self, event):
+        if self.debug:
+            print("Pressed", event.key)
+
+        ACTION_MAPPING = {
+            "left": Actions.left,
+            "right": Actions.right,
+            "up": Actions.up,
+            "down": Actions.down,
+            " ": Actions.interact,
+            "tab": Actions.stay,
+        }
+
+        match event.key:
+            case "escape":
+                self.viz.window.close()
+                return
+            case "backspace":
+                self._reset()
+                return
+            case key if key in ACTION_MAPPING:
+                action = ACTION_MAPPING[key]
+            case key:
+                print(f"Key {key} not recognized")
+                return
+
+        self._step(action)
+
+    def _redraw(self):
+        self.viz.render(self.state, agent_view_size=self.env.agent_view_size)
+
+    def _reset(self):
+        self.key, key = jax.random.split(self.key)
+        _, state = jax.jit(self.env.reset)(key)
+        self.state = state
+
+        self._redraw()
+
+    def _step(self, action):
+        self.key, subkey = jax.random.split(self.key)
+
+        actions = {f"agent_{i}": jnp.array(action) for i in range(self.env.num_agents)}
+        if self.debug:
+            print("Actions: ", actions)
+
+        obs, state, reward, done, info = jax.jit(self.env.step_env)(
+            subkey, self.state, actions
+        )
+        self.state = state
+        print(f"t={state.time}: reward={reward['agent_0']}, done = {done['__all__']}")
+
+        if self.debug:
+            a0_obs = obs["agent_0"]
+            a0_obs = jnp.transpose(a0_obs, (2, 0, 1))
+            print("Agent 0 observation: ", a0_obs)
+            print("Reward: ", reward)
+            print("Shaped reward: ", info["shaped_reward"])
+
+        if done["__all__"]:
+            self._reset()
+        else:
+            self._redraw()
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--layout", type=str, help="Overcooked layout", default="cramped_room"
+    )
+    parser.add_argument(
+        "--seed",
+        type=int,
+        help="random seed to generate the environment with",
+        default=0,
+    )
+    parser.add_argument(
+        "--agent_view_size",
+        type=int,
+        help="Number of cells in the agent view. If not provided, the agent will see the whole grid.",
+    )
+    parser.add_argument(
+        "--no_jit",
+        default=False,
+        help="Disable JIT compilation",
+        action="store_true",
+    )
+    parser.add_argument(
+        "--debug", default=False, help="Debug mode", action="store_true"
+    )
+    args = parser.parse_args()
+
+    if len(args.layout) == 0:
+        raise ValueError("You must provide a layout.")
+    layout = layouts[args.layout]
+
+    interactive = InteractiveOvercookedV3(
+        layout=layout,
+        agent_view_size=args.agent_view_size,
+        no_jit=args.no_jit,
+        debug=args.debug,
+    )
+
+    key = jax.random.PRNGKey(args.seed)
+    interactive.run(key)