Add collisions as part of the simulator state and change default values

requirements.txt

@@ -10,6 +10,7 @@ matplotlib==3.8.2
 
 # Client-server communication
 grpcio==1.60.0
+grpcio-tools==1.62.1
 
 # Tests and code formatting
 pytest==8.0.0

scripts/run_server.py

@@ -14,6 +14,9 @@
 lg = logging.getLogger(__name__)
 
 def parse_args():
+    # TODO : Quick nnote : would be cool to have a separate file to have all the default values 
+    # In fact the best solution might be to use hydra with yaml files instead of argparse, 
+    # Would also be pretty convenient to store positions etcc for different notebook scenes
     parser = argparse.ArgumentParser(description='Simulator Configuration')
     parser.add_argument('--box_size', type=float, default=100.0, help='Size of the simulation box')
     parser.add_argument('--n_agents', type=int, default=10, help='Number of agents')
@@ -26,6 +29,8 @@ def parse_args():
     parser.add_argument('--to_jit', action='store_false', help='Whether to use JIT compilation')
     parser.add_argument('--use_fori_loop', action='store_true', help='Whether to use fori loop')
     parser.add_argument('--log_level', type=str, default='INFO', help='Logging level')
+    parser.add_argument('--col_eps', type=float, required=False, default=0.003)
+    parser.add_argument('--col_alpha', type=float, required=False, default=0.7)
 
     return parser.parse_args()
 
@@ -44,7 +49,9 @@ def parse_args():
         freq=args.freq,
         neighbor_radius=args.neighbor_radius,
         to_jit=args.to_jit,
-        use_fori_loop=args.use_fori_loop
+        use_fori_loop=args.use_fori_loop,
+        col_eps=args.col_eps,
+        col_alpha=args.col_alpha
     )
 
     agent_configs = [AgentConfig(idx=i,
@@ -65,5 +72,6 @@ def parse_args():
                                         })
 
     simulator = Simulator(state, behaviors.behavior_bank, dynamics_rigid)
+
     lg.info('Simulator server started')
     serve(simulator)

scripts/run_simulation.py

@@ -13,6 +13,7 @@
 
 def parse_args():
     parser = argparse.ArgumentParser(description='Simulator Configuration')
+    parser.add_argument('--log_level', type=str, default='INFO', help='Logging level')
     # Experiment run arguments
     parser.add_argument('--num_loops', type=int, default=10, help='Number of simulation loops')
     # Simulator config arguments
@@ -26,7 +27,8 @@ def parse_args():
     # By default jit compile the code and use normal python loops
     parser.add_argument('--to_jit', action='store_false', help='Whether to use JIT compilation')
     parser.add_argument('--use_fori_loop', action='store_true', help='Whether to use fori loop')
-    parser.add_argument('--log_level', type=str, default='INFO', help='Logging level')
+    parser.add_argument('--col_eps', type=float, required=False, default=0.003)
+    parser.add_argument('--col_alpha', type=float, required=False, default=0.7)
 
     return parser.parse_args()
 
@@ -45,7 +47,9 @@ def parse_args():
         freq=args.freq,
         neighbor_radius=args.neighbor_radius,
         to_jit=args.to_jit,
-        use_fori_loop=args.use_fori_loop
+        use_fori_loop=args.use_fori_loop,
+        col_eps=args.col_eps,
+        col_alpha=args.col_alpha
     )
 
     agent_configs = [

vivarium/controllers/config.py

@@ -83,6 +83,8 @@ class SimulatorConfig(Config):
     neighbor_radius = param.Number(100., bounds=(0, None))
     to_jit = param.Boolean(True)
     use_fori_loop = param.Boolean(False)
+    col_eps = param.Number(0.003)
+    col_alpha = param.Number(0.7)
 
     def __init__(self, **params):
         super().__init__(**params)

vivarium/controllers/converters.py

@@ -25,7 +25,7 @@
 object_state_fields = [f.name for f in jax_md.dataclasses.fields(ObjectState)]
 
 object_common_fields = [f for f in object_config_fields if f in object_state_fields]
-#
+
 simulator_config_fields = SimulatorConfig.param.objects().keys()
 simulator_state_fields = [f.name for f in jax_md.dataclasses.fields(SimulatorState)]
 
@@ -106,7 +106,9 @@ def get_default_state(n_entities_dict):
                                                 n_agents=jnp.array([n_agents]), n_objects=jnp.array([n_objects]),
                                                 num_steps_lax=jnp.array([1]), dt=jnp.array([1.]), freq=jnp.array([1.]),
                                                 neighbor_radius=jnp.array([1.]),
-                                                to_jit= jnp.array([1]), use_fori_loop=jnp.array([0])),
+                                                to_jit= jnp.array([1]), use_fori_loop=jnp.array([0]),
+                                                col_alpha=jnp.array([0.]),
+                                                col_eps=jnp.array([0.])),
                  nve_state=NVEState(position=RigidBody(center=jnp.zeros((n_entities, 2)), orientation=jnp.zeros(n_entities)),
                                     momentum=None,
                                     force=RigidBody(center=jnp.zeros((n_entities, 2)), orientation=jnp.zeros(n_entities)),

vivarium/simulator/grpc_server/converters.py

@@ -22,7 +22,9 @@ def proto_to_simulator_state(simulator_state):
                           freq=proto_to_ndarray(simulator_state.freq),
                           neighbor_radius=proto_to_ndarray(simulator_state.neighbor_radius),
                           to_jit=proto_to_ndarray(simulator_state.to_jit),
-                          use_fori_loop=proto_to_ndarray(simulator_state.use_fori_loop)
+                          use_fori_loop=proto_to_ndarray(simulator_state.use_fori_loop),
+                          col_eps=proto_to_ndarray(simulator_state.col_eps),
+                          col_alpha=proto_to_ndarray(simulator_state.col_alpha)
                           )
 
 
@@ -81,7 +83,9 @@ def simulator_state_to_proto(simulator_state):
         freq=ndarray_to_proto(simulator_state.freq),
         neighbor_radius=ndarray_to_proto(simulator_state.neighbor_radius),
         to_jit=ndarray_to_proto(simulator_state.to_jit),
-        use_fori_loop=ndarray_to_proto(simulator_state.use_fori_loop)
+        use_fori_loop=ndarray_to_proto(simulator_state.use_fori_loop),
+        col_eps=ndarray_to_proto(simulator_state.col_eps),
+        col_alpha=ndarray_to_proto(simulator_state.col_alpha)
     )
 
 

vivarium/simulator/grpc_server/protos/simulator.proto

@@ -66,6 +66,8 @@ message SimulatorState {
   NDArray neighbor_radius = 8;
   NDArray to_jit = 9;
   NDArray use_fori_loop = 10;
+  NDArray col_eps = 11;
+  NDArray col_alpha = 12;
 }
 
 message NVEState {