Add interaction between server, vis client and notebook client

client.py

@@ -6,7 +6,7 @@
 
 from simulationsandbox.two_d_simulation import SimpleSimulation
 from simulationsandbox.utils.network import SERVER
-from simulationsandbox.sim_types import SIMULATIONS
+from simulationsandbox.utils.sim_types import SIMULATIONS
 
 
 PORT = 5050

notebook_controller.ipynb

@@ -21,8 +21,8 @@
     "import jax.numpy as jnp\n",
     "from flax import serialization\n",
     "\n",
-    "from MultiAgentsSim.two_d_simulation import SimpleSimulation\n",
-    "from MultiAgentsSim.utils.network import SERVER"
+    "from simulationsandbox.two_d_simulation import SimpleSimulation\n",
+    "from simulationsandbox.utils.network import SERVER"
    ]
   },
   {
@@ -89,7 +89,16 @@
     "    current_state = serialization.from_bytes(state_example, client.recv(state_bytes_size))\n",
     "    response_state = change_state_agent_color(current_state, agent_idx, color)\n",
     "    client.send(serialization.to_bytes(response_state))\n",
-    "    return "
+    "    return \n",
+    "\n",
+    "def pause(client):\n",
+    "    client.send(\"PAUSE\".encode())\n",
+    "\n",
+    "def resume(client):\n",
+    "    client.send(\"RESUME\".encode())\n",
+    "\n",
+    "def stop(client):\n",
+    "    client.send(\"STOP\".encode())"
    ]
   },
   {

server.py

@@ -9,7 +9,7 @@
 
 from simulationsandbox.wrapper import SimulationWrapper
 from simulationsandbox.utils.network import SERVER
-from simulationsandbox.sim_types import SIMULATIONS
+from simulationsandbox.utils.sim_types import SIMULATIONS
 
 parser = argparse.ArgumentParser()
 parser.add_argument('--sim_type', type=str, default="two_d")
@@ -29,6 +29,7 @@
 MAX_AGENTS = 10
 NUM_OBS = 3 
 GRID_SIZE = 20 
+PRINT_DATA = False
 
 # Initialize server
 server = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
@@ -50,7 +51,7 @@
 update_event = threading.Event()
 
 
-simulation = SimulationWrapper(sim, state, key, step_delay=STEP_DELAY, update_event=update_event)
+simulation = SimulationWrapper(sim, state, key, step_delay=STEP_DELAY, update_event=update_event, print_data=PRINT_DATA)
 print(f"{len(pickle.dumps(simulation.state))}")
 
 
@@ -102,10 +103,32 @@ def communicate_with_client(client, addr, connection_type):
 
                     elif request == "SET_STATE":
                         with sim_lock:
+                            simulation.pause()
                             client.send(serialization.to_bytes(simulation.state))  
                             updated_state = serialization.from_bytes(state, client.recv(state_byte_size))  
                             simulation.state = updated_state
+                            simulation.resume()
+
+                    elif request == "PAUSE":
+                        with sim_lock:
+                            simulation.pause()
+                        print("Simulation paused")
 
+                    elif request == "RESUME":
+                        with sim_lock:
+                            simulation.resume()
+                        print("Simulation resumed")
+
+                    elif request == "STOP":
+                        with sim_lock:
+                            simulation.stop()
+                        print("Simulation stopped")
+
+                    elif request == "START":
+                        with sim_lock:
+                            simulation.start()
+                        print("Simulation started")
+
                     else:
                         print(f"Unknow request type {request}")
 

simulationsandbox/utils/network.py

@@ -4,4 +4,19 @@
 servers = ['localhost', '10.204.2.210', '10.204.2.229', '10.204.2.189', '192.168.1.24']
 server_idx = 0
 
-SERVER = servers[server_idx]
+SERVER = servers[server_idx]
+
+
+# Establish a connection with a client
+def establish_connection(client, addr, simulation, sim_type, data_size):
+    try:
+        client.send(sim_type.encode())
+        client.send(pickle.dumps(simulation.state))
+        connection_type = client.recv(data_size).decode()
+        print(f"{connection_type} connection established with {addr}")
+        return connection_type
+
+    except socket.error as e:
+        print(f"error: {e}")
+        client.close()
+        print(f"Client {client} disconnected")