Multiplayer over the network

.gitignore

@@ -1,4 +1,5 @@
 /venv/
 **/__pycache__/
 *-env
-robotouille.egg-info
+robotouille.egg-info
+recordings

Dockerfile

@@ -0,0 +1,7 @@
+FROM python:3.9.6
+
+COPY . .
+
+RUN pip install -r requirements.txt
+
+CMD python main.py --environment_name "cook_soup" --role server

main.py

@@ -4,7 +4,11 @@
 parser = argparse.ArgumentParser()
 parser.add_argument("--environment_name", help="The name of the environment to create.", default="original")
 parser.add_argument("--seed", help="The seed to use for the environment.", default=None)
+parser.add_argument("--role", help="\"simulator\" for vanilla simulator, \"client\" if client, \"server\" if server, \"single\" if single-player, \"replay\" if replaying, \"render\" if rendering video", default="simulator")
+parser.add_argument("--server_display", action="store_true", help="Whether to show the game window as server (ignored for other roles)")
+parser.add_argument("--host", help="Host to connect to", default="ws://localhost:8765")
+parser.add_argument("--replay", help="Recording to replay", default="")
 parser.add_argument("--noisy_randomization", action="store_true", help="Whether to use 'noisy randomization' for procedural generation")
 args = parser.parse_args()
 
-simulator(args.environment_name, args.seed, args.noisy_randomization)
+simulator(args.environment_name, args.seed, args.role, args.server_display, args.host, args.replay, args.noisy_randomization)

networking/README.md

@@ -0,0 +1,40 @@
+# Robotouille Networking README
+
+## Overview
+
+Built into Robotouille is networking support for multiplayer. Robotouille uses an authoritative server, which clients can connect to. Servers record games played on them to collect data. 
+
+## Available Modes
+
+Several modes are offered, which can be chosen using the `--role` argument:
+
+1. `server`
+2. `client `
+3. `single`
+4. `replay`
+5. `render`
+6. `simulator`
+
+## Simulator
+
+This mode runs Robotouille without any networking overhead.
+
+## Server
+
+This mode sets up an Robotouille server. Clients that connect are automatically matchmaked and put together into lobbies. Use argument `display_server` to have the server render active games.
+
+## Client
+
+This mode runs the Robotouille client. Use argument `host` to choose which host to connect to. Defaults to local host.
+
+## Single
+
+This mode runs both the server and client for a single player experience. Server features, such as game recordings, remain available.
+
+## Replay
+
+This mode replays a recorded game through a window. The recording is specified with argument `recording`.
+
+## Render
+
+This mode renders a recording game into a video. The video is exported to the recordings folder. The recording is specified with argument `recording`.

networking/client.py

@@ -0,0 +1,64 @@
+import asyncio
+import json
+import pickle
+import base64
+import websockets
+import pygame
+from robotouille.robotouille_env import create_robotouille_env
+from utils.robotouille_input import create_action_from_control
+
+def run_client(environment_name: str, seed: int = 42, host: str="ws://localhost:8765", noisy_randomization: bool = False):
+    asyncio.run(client_loop(environment_name, seed, host, noisy_randomization))
+
+async def client_loop(environment_name: str, seed: int = 42, host: str="ws://localhost:8765", noisy_randomization: bool = False):
+    uri = host
+
+    async def send_actions(websocket, shared_state):
+        env = shared_state["env"]
+        renderer = shared_state["renderer"]
+        renderer.render(shared_state["obs"], mode='human')
+        player = shared_state["player"]
+        online = True
+        while not shared_state["done"]:
+            pygame_events = pygame.event.get()
+            mousedown_events = list(filter(lambda e: e.type == pygame.MOUSEBUTTONDOWN, pygame_events))
+            keydown_events = list(filter(lambda e: e.type == pygame.KEYDOWN, pygame_events))
+            action, args = create_action_from_control(env, shared_state["obs"], player, mousedown_events + keydown_events, renderer)
+
+            # Use this to simulate disconnect
+            online = not (pygame.key.get_mods() & pygame.KMOD_CAPS)
+
+            if action is not None:
+                if online:
+                    encoded_action = base64.b64encode(pickle.dumps((action, args))).decode('utf-8')
+                    await websocket.send(json.dumps(encoded_action))
+            renderer.render(env.get_state(), mode='human')
+
+            await asyncio.sleep(0)  # Yield control to allow other tasks to run
+
+    async def receive_responses(websocket, shared_state):
+        while not shared_state["done"]:
+            response = await websocket.recv()
+            data = json.loads(response)
+            shared_state["done"] = data["done"]
+
+            shared_state["env"].set_state(pickle.loads(base64.b64decode(data["env"])))
+            shared_state["obs"] = pickle.loads(base64.b64decode(data["obs"]))
+
+    async with websockets.connect(uri) as websocket:
+        env, _, renderer = create_robotouille_env(environment_name, seed, noisy_randomization)
+        obs, info = env.reset()
+        shared_state = {"done": False, "env": env, "renderer": renderer, "obs": obs, "player": None}
+        print("In lobby")
+
+        opening_message = await websocket.recv()
+        print("In game")
+        opening_data = json.loads(opening_message)
+        player_index = pickle.loads(base64.b64decode(opening_data["player"]))
+        player = env.get_state().get_players()[player_index]
+        shared_state["player"] = player
+
+        sender = asyncio.create_task(send_actions(websocket, shared_state))
+        receiver = asyncio.create_task(receive_responses(websocket, shared_state))
+        await asyncio.gather(sender, receiver)
+        # Additional cleanup if necessary

networking/server.py

@@ -0,0 +1,128 @@
+import asyncio
+import json
+import pickle
+import base64
+import websockets
+import time
+from pathlib import Path
+from datetime import datetime
+from robotouille.robotouille_env import create_robotouille_env
+
+SIMULATE_LATENCY = False
+SIMULATED_LATENCY_DURATION = 0.25
+
+def run_server(environment_name: str, seed: int=42, noisy_randomization: bool=False, display_server: bool=False, event: asyncio.Event=None):
+    asyncio.run(server_loop(environment_name, seed, noisy_randomization, display_server))
+
+async def server_loop(environment_name: str, seed: int=42, noisy_randomization: bool=False, display_server: bool=False, event: asyncio.Event=None):
+    waiting_queue = {}
+    reference_env = create_robotouille_env(environment_name, seed, noisy_randomization)[0]
+    num_players = len(reference_env.get_state().get_players())
+
+    async def simulator(connections):
+        try:
+            print("Start game", connections.keys())
+            recording = {}
+            recording["start_time"] = datetime.now().strftime("%Y%m%d_%H%M%S_%f")
+            recording["environment_name"] = environment_name
+            recording["seed"] = seed
+            recording["noisy_randomization"] = noisy_randomization
+            recording["actions"] = []
+            recording["violations"] = []
+            start_time = time.monotonic()
+
+            env, json_data, renderer = create_robotouille_env(environment_name, seed, noisy_randomization)
+            obs, info = env.reset()
+            if display_server:
+                renderer.render(obs, mode='human')
+            done = False
+            interactive = False  # Adjust based on client commands later if needed
+
+            assert len(connections) == num_players
+            sockets_to_playerID = {}
+            for i, socket in enumerate(connections.keys()):
+                sockets_to_playerID[socket] = i
+                player_data = base64.b64encode(pickle.dumps(i)).decode('utf-8')
+                opening_message = json.dumps({"player": player_data})
+                await socket.send(opening_message)
+
+            while not done:
+                # Wait for messages from any client
+                # TODO(aac77): make more robust
+                receive_tasks = {asyncio.create_task(q.get()): client for client, q in connections.items()}
+                finished_tasks, pending_tasks = await asyncio.wait(receive_tasks.keys(), return_when=asyncio.FIRST_COMPLETED)
+
+                # Cancel pending tasks, otherwise we leak
+                for task in pending_tasks:
+                    task.cancel()
+
+                # Retrieve the message from the completed task
+                actions = [(None, None)] * num_players
+                for task in finished_tasks:
+                    message = task.result()
+                    client = receive_tasks[task]
+                    encoded_action = json.loads(message)
+                    action = pickle.loads(base64.b64decode(encoded_action))
+
+                    actions[sockets_to_playerID[client]] = action
+
+
+                if SIMULATE_LATENCY:
+                    time.sleep(SIMULATED_LATENCY_DURATION)
+
+                reply = None
+
+                try:
+                    obs, reward, done, info = env.step(actions, interactive=interactive)
+                    recording["actions"].append((actions, env.get_state(), time.monotonic() - start_time))
+                    if display_server:
+                        renderer.render(obs, mode='human')
+                except AssertionError:
+                    print("violation")
+                    recording["violations"].append((actions, time.monotonic() - start_time))
+
+                env_data = pickle.dumps(env.get_state())
+                encoded_env_data = base64.b64encode(env_data).decode('utf-8')
+                obs_data = pickle.dumps(obs)
+                encoded_obs_data = base64.b64encode(obs_data).decode('utf-8')
+                reply = json.dumps({"env": encoded_env_data, "obs": encoded_obs_data, "done": done})
+
+                if SIMULATE_LATENCY:
+                    time.sleep(SIMULATED_LATENCY_DURATION)
+                websockets.broadcast(connections.keys(), reply)
+            recording["result"] = "done"
+        except BaseException as e:
+            traceback.print_exc(e)
+            recording["result"] = traceback.format_exc(e)
+        finally:
+            for websocket in connections.keys():
+                await websocket.close()
+
+            if display_server:
+                renderer.render(obs, close=True)
+            print("GG")
+            recording["end_time"] = datetime.now().strftime("%Y%m%d_%H%M%S_%f")
+
+            p = Path('recordings')
+            p.mkdir(exist_ok=True)
+            with open(p / (recording["start_time"] + '.pkl'), 'wb') as f:
+                pickle.dump(recording, f)
+
+    async def handle_connection(websocket):
+        # TODO(aac77): make more robust
+        print("Hello client", websocket)
+        q = asyncio.Queue()
+        waiting_queue[websocket] = q
+        if len(waiting_queue) == num_players:
+            connections = waiting_queue.copy()
+            waiting_queue.clear()
+            asyncio.create_task(simulator(connections))
+        async for message in websocket:
+            await q.put(message)
+
+    if event == None:
+        event = asyncio.Event()
+
+    async with websockets.serve(handle_connection, "0.0.0.0", 8765):
+        print("I am server")
+        await event.wait()

requirements.txt

@@ -5,6 +5,7 @@ fonttools==4.38.0
 gym==0.26.2
 gym-notices==0.0.8
 imageio==2.25.1
+imageio-ffmpeg==0.4.9
 importlib-metadata==6.0.0
 importlib-resources==5.12.0
 kiwisolver==1.4.4
@@ -21,4 +22,5 @@ scikit-image==0.19.3
 scipy==1.10.0
 six==1.16.0
 tifffile==2023.2.3
+websockets==12.0
 zipp==3.14.0

robotouille/env.py

@@ -302,7 +302,7 @@ def build_state(domain_json, environment_json):
     true_predicates += build_stacking_predicates(environment_json)
     goal = build_goal(environment_json)
 
-    state = State().initialize(domain, objects, true_predicates, goal)
+    state = State().initialize(domain, objects, true_predicates, goal, [])
 
     return state
 

robotouille/robotouille_simulator.py

@@ -1,9 +1,97 @@
 import pygame
 from utils.robotouille_input import create_action_from_control
 from robotouille.robotouille_env import create_robotouille_env
+import asyncio
+import json
+import pickle
+import base64
+import websockets
+import time
+from pathlib import Path
+from datetime import datetime
+import imageio
+import traceback
+import networking.server as robotouille_server
+import networking.client as robotouille_client
 
+def simulator(environment_name: str, seed: int=42, role: str="simulator", display_server: bool=False, host: str="ws://localhost:8765", recording: str="", noisy_randomization: bool=False):
+    if recording != "" and role != "replay" and role != "render":
+        role = "replay"
+    if role == "server":
+        robotouille_server.run_server(environment_name, seed, noisy_randomization, display_server)
+    elif role == "client":
+        robotouille_client.run_client(environment_name, seed, host, noisy_randomization)
+    elif role == "single":
+        asyncio.run(single_player(environment_name, seed, noisy_randomization))
+    elif role == "replay":
+        replay(recording)
+    elif role == "render":
+        render(recording)
+    elif role == "simulator":
+        simulate(environment_name, seed, noisy_randomization)
+    else:
+        print("Invalid role:", role)
 
-def simulator(environment_name: str, seed: int=42, noisy_randomization: bool=False):
+async def single_player(environment_name: str, seed: int=42, noisy_randomization: bool=False):
+    event = asyncio.Event()
+    server = asyncio.create_task(robotouille_server.server_loop(environment_name=environment_name, seed=seed, noisy_randomization=noisy_randomization, event=event))
+    await asyncio.sleep(0.5)  # wait for server to initialize
+    client = asyncio.create_task(robotouille_client.client_loop(environment_name=environment_name, seed=seed, noisy_randomization=noisy_randomization))
+    await client
+    event.set()
+    await server
+
+def replay(recording_name: str):
+    if not recording_name:
+        raise ValueError("Empty recording_name supplied")
+
+    p = Path('recordings')
+    with open(p / (recording_name + '.pkl'), 'rb') as f:
+        recording = pickle.load(f)
+
+    env, _, renderer = create_robotouille_env(recording["environment_name"], recording["seed"], recording["noisy_randomization"])
+    obs, _ = env.reset()
+    renderer.render(obs, mode='human')
+
+    previous_time = 0
+    for actions, state, t in recording["actions"]:
+        time.sleep(t - previous_time)
+        previous_time = t
+        obs, reward, done, info = env.step(actions=actions, interactive=False)
+        renderer.render(obs, mode='human')
+    renderer.render(obs, close=True)
+
+def render(recording_name: str):
+    p = Path('recordings')
+    with open(p / (recording_name + '.pkl'), 'rb') as f:
+        recording = pickle.load(f)
+
+    env, _, renderer = create_robotouille_env(recording["environment_name"], recording["seed"], recording["noisy_randomization"])
+    obs, _ = env.reset()
+    frame = renderer.render(obs, mode='rgb_array')
+
+    vp = Path('recordings')
+    vp.mkdir(exist_ok=True)
+    fps = 20
+    video_writer = imageio.get_writer(vp / (recording_name + '.mp4'), fps=fps)
+
+    i = 0
+    t = 0
+    while i < len(recording["actions"]):
+        actions, state, time_stamp = recording["actions"][i]
+        while t > time_stamp:
+            obs, reward, done, info = env.step(actions=actions, interactive=False)
+            frame = renderer.render(obs, mode='rgb_array')
+            i += 1
+            if i >= len(recording["actions"]):
+                break
+            action, state, time_stamp = recording["actions"][i]
+        t += 1 / fps
+        video_writer.append_data(frame)
+    renderer.render(obs, close=True)
+    video_writer.close()
+
+def simulate(environment_name, seed, noisy_randomization):
     # Your code for robotouille goes here
     env, json, renderer = create_robotouille_env(environment_name, seed, noisy_randomization)
     obs, info = env.reset()
@@ -30,4 +118,4 @@ def simulator(environment_name: str, seed: int=42, noisy_randomization: bool=Fal
             continue
         obs, reward, done, info = env.step(actions, interactive=interactive)
         renderer.render(obs, mode='human')
-    renderer.render(obs, close=True)
+    renderer.render(obs, close=True)