Init state without configs #57

corentinlger · 2024-03-20T16:01:21Z

Description

The PR modifies a lot of lines in the code, but there aren't tricky logic changes, we just modify the way to initialize the state without configs. I'll try to be as clear as I can when explaining the modifications.

I moved all the state classes and definitions from sim_computation.py to a new file called states.py.

The modifications also include refactoring imports in multiple files to align with the new file structure, enhancing readability and preventing linting errors.

In states.py, five functions have been implemented to initialize various states, such as simulator_state, nve_state, agent_state, objects_state, and the combined state. These changes offer greater flexibility in initializing states, allowing for customization such as specifying agent positions or using random values. Future iterations will extend this approach to more attributes, such as colors, diameters, behaviors ... to provide even more customization options.

Additionally, default values have been defined within these classes, and type hints have been added to functions for clarity. But I can still add some more docstrings to define what each parameter does.

Related Issue (if applicable)

closes #51

How to Test

Launch the server

python3 scripts/run_server.py

Launch the Panel interface

panel serve scripts/run_interface.py --autoreload

Check if nothing is broken with these two lines, normally it should be the case because only code changing the state init has been modified.

Normally I added some ugly lines to show that you can initiate the position of agents with a list in run_server. Run :

python3 scripts/run_server.py --custom_pos

You can also try to launch the simulation with a custom number of agents or objects (will be usefull later for eco-evo simulations). Actually for this feature I think it would make more sense to rename the current n_agents to max_agents, and n_existing_agents to n_agents (same for objects).

python3 scripts/run_server.py --n_existing_agents 7 --n_existing_objects 1

And relaunch the interface, check that the agents are indeed defined on the diagonal of the box. Later, this will easiely enable to define custom scenes, notably when using external files to store th epositions, colors ... as mentionned in #54

Screenshots (if applicable)

…corentin/init_state_without_configs

Marsolo1 · 2024-03-25T13:00:32Z

scripts/run_server.py

+    else:
+        positions = None
+
+    simulator_state = init_simulator_state(


This part works to show that we can use a list of coordinates, but I think it should be removed before the merge

Indeed, this is meant to be deleted. I just added it to make the testing of the PR easier

Marsolo1

The new system seems to work fine for me, when this PR is merged i'll add the Hydra config files so that we can finally have scene files

clement-moulin-frier

I have started to review this PR but unfortunately I don't have time to finish this morning. Please start to have a look at my comments so that we can discuss it this afternoon. I'll try to finish it today.

Merge last master(no collision parameters)
When running the server with --n_existing_agents 7 --n_existing_objects 1, then selecting a non-existing agent and object and activating the the visible checkbox, the selected entity appear for a very short time and disappear right away. I'm not sure it is specific to this PR though, it might be a more general bug. Can you have a look @Marsolo1 ?

clement-moulin-frier · 2024-03-27T14:26:34Z

scripts/run_server.py

 from vivarium.simulator.grpc_server.simulator_server import serve

 lg = logging.getLogger(__name__)

 def parse_args():
    parser = argparse.ArgumentParser(description='Simulator Configuration')
+    parser.add_argument('--custom_pos', action='store_true', help='Just a test arg to wether use custom or random pos')


Check if you want to keep it in the merge. A typo in wether

EDIT: I see in your reply to Martial below that you plan to remove it, all good

clement-moulin-frier · 2024-03-28T10:05:11Z

scripts/run_server.py

+    parser.add_argument('--n_existing_agents', type=int, default=10, help='Number of agents')
    parser.add_argument('--n_objects', type=int, default=2, help='Number of objects')
+    parser.add_argument('--n_existing_objects', type=int, default=2, help='Number of agents')


I agree with your proposition that " it would make more sense to rename the current n_agents to max_agents, and n_existing_agents to n_agents (same for objects)."

Not sure if you already do it, but it will be useful to add an assert checking that the number of existing agents/objects does not exceed the max number

Yes I already check this

clement-moulin-frier · 2024-03-28T10:10:11Z

tests/test_simulator_init.py

+          n_objects=n_objects)
+
+    nve_state = init_nve_state(
+        simulator_state,


wrong indentation

clement-moulin-frier · 2024-03-28T10:11:10Z

tests/test_simulator_init.py

+
+
+def test_init_simulator_args():
+    """ Test the initialization of state with arguments """


You might also want to test the case where we only provide a subset of the arguments, to see if it merges well with the default values

Another useful test to do would be to check if (some of) the values in the created state are indeed the ones that are provided in the args or defaults (e.g. just adding a few assert in all your test functions)

Indeed I'll add that

clement-moulin-frier · 2024-03-28T10:26:00Z

vivarium/simulator/states.py

+from typing import Optional, List, Union
+from enum import Enum
+
+import matplotlib.colors as mcolors


Not related to this PR, but I just notice that this requires matplotlib. We should either:

Add matplotlib to the requirements.txt (if not already there)

Or: find a way to do the same thing without using matlplotlib. The mcolors function is I think just used to convert colors between an vector RGB format (used in State) and a string format (used in the config). Maybe there is a way to do this without relying on matlplotlib.

If possible, I think the second option is better, because imposing to install matlplotlib (which is quite a big lib) just for this. We can discuss it and do it later though, you can just create a new issue atm

There is actually a deeper issue here, see my other comment below

I agree with you only using matplotlib for this isn't really useful. But it is needed to launch the interface with panel anyway I am not wrong. That's why it was already in the requirements.txt.

clement-moulin-frier · 2024-03-28T10:32:52Z

vivarium/simulator/states.py

+    return jnp.array(list(mcolors.to_rgb(color_str)))
+
+
+def init_simulator_state(


We can do this in another Issue/PR, but it would be better if default values are defined at a single place in the code. For now we have them in several places, e.g. here, in the run scripts, in the Config classes (and maybe elsewhere), this will be very hard to maintain. We could e.g. have them in a dedicated file and all the classes/functions that need them would import them from there. Might be easier to do it in the context of #30 though (if you agree please update #30 or related PR so that we don't forget, otherwise let's discuss it this afternoon)

clement-moulin-frier · 2024-03-28T10:34:09Z

vivarium/simulator/states.py

+
+
+def _init_positions(key_pos, positions, n_elements, box_size, n_dims=2):
+    assert (len(positions) == n_elements if positions else True)


To check but assert (positions is None or len(positions) == n_elements) is more clear ;)

clement-moulin-frier · 2024-03-28T10:38:00Z

vivarium/simulator/states.py

+        positions = random.uniform(key_pos, (n_elements, n_dims)) * box_size
+    return positions
+
+def _init_existing(n_existing, n_elements):


I am not sure what an "element" refers to here. Isn't it instead related to an entity? In that case change it to e.g. n_entities, it will be more clear (in general we should avoid multiplying the terms we use)

If you do it report it everywhere you use '"element"

clement-moulin-frier · 2024-03-28T10:40:30Z

vivarium/simulator/states.py

+# Also think it would be easier to only handle the nve state in the physics engine. (if it doesn't make the  simulation run slower)
+# Here it makes a really long function, and it isn't very modular 
+# Plus we store different attributes of entities (e.g Agents positions and sensors) in different dataclasses


Let's discuss it, as we said last time it is part of a more general issue about how we handle conversions between data types that encode the simular infomation (AgentState, NVEState, Config, proto ..)

You can create an issue about explaining what we have already discussed on this topic (a potential solution would be to have useful helper functions in the State class)

Yes I wrote these comments before our discussion, I'll create an issue

clement-moulin-frier · 2024-03-28T10:45:46Z

vivarium/simulator/states.py

+    n_objects = simulator_state.n_objects[0]
+    n_entities = n_agents + n_objects
+
+    key = random.PRNGKey(seed)


I'm not sure this follows the good practice about how to deal with randomness in JAX, see https://jax.readthedocs.io/en/latest/jax-101/05-random-numbers.html

Typically, the seed should be defined in a single place in the code, and all functions requiring a random number should use jax.random.split. PRNGKey should be called only once in the entire code base. (it's how I understand it at least)

See if you can fix it here, otherwise create an issue (point at the code here so that we don't forget it)

It is a good question. I am not sure about this (because in our case this key only serves to initialize random positions, and they key in simulator serve other purposes) but it will be easy to change if we want to. Actually I checked some RL repos in Jax and they seem to define keys with seeds in different parts of the code (ex: in environments and in algorithms).

From what I saw we only have another key initialization in state, we can just change the init methods so it doesn't accept a seed that generates a key, but directly a key (that you generate at the beginning of the script you run).

Marsolo1 · 2024-03-28T13:22:47Z

When running the server with `--n_existing_agents 7 --n_existing_objects 1`, then selecting a non-existing agent and object and activating the the `visible` checkbox, the selected entity appear for a very short time and disappear right away. I'm not sure it is specific to this PR though, it might be a more general bug. Can you have a look @Marsolo1 ?

Yes I already saw this bug and fixed it but it was in another PR, so I just pushed the fix to main directly (it was only a False in the wrong place). If we merge this branch with main it should fix it here too

clement-moulin-frier

This is the second part of my review, which is now completed. Please read the comments and then we discuss

clement-moulin-frier · 2024-03-28T13:50:08Z

vivarium/simulator/states.py

+    key_pos, key_or = random.split(key)
+    key_ag, key_obj = random.split(key_pos)


Not very standard, see JAX doc on split is usually used

you prefer to do a random.split(key, 4) ?

clement-moulin-frier · 2024-03-28T13:57:08Z

vivarium/simulator/states.py

+        use_fori_loop=jnp.array([1*use_fori_loop]))
+
+
+def _init_positions(key_pos, positions, n_elements, box_size, n_dims=2):


Use the the constant SPACE_NDIMS defined in sim_computation.

Also I'm not sure that having a single function for either casting positions into a jax array, or generating random ones, helps. It makes the function hard quite strange (if one wants to just cast existing positions, most of this function's arguments become useless), whereas casting to JAX array is just a single line. I would recommend to simply have a function for generating random positions ; and just explicitly cast to jnp.array where it is needed without calling this function

clement-moulin-frier · 2024-03-28T13:58:40Z

vivarium/simulator/states.py

+    existing_objects = _init_existing(existing_objects, n_objects)
+    exists = jnp.concatenate((existing_agents, existing_objects), dtype=int)
+
+    # TODO: Why is momentum set to none ?


This is the way it is done in jax_md documentation I think

then jax_md initialize it (in our step function I think)

clement-moulin-frier · 2024-03-28T14:01:06Z

vivarium/simulator/states.py

+        diameter: float = 5.,
+        friction: float = 0.1,
+        mass_center: float = 1.,
+        mass_orientation: float = 0.125,


We should be able to pass a full array as well here, e.g. if we want to set different diameters for different entities. A solution could be to check the type of this argument : if it is a scalar, all entities will have the same value, if it is an array, we set them as such. Will require to explain it in a docstring ;)

Indeed I indicated this in a comment in the file normally, I just didn't do it for every attribute

clement-moulin-frier · 2024-03-28T14:02:18Z

vivarium/simulator/states.py

+        )
+
+
+# Could implement it as a static or class method


It's true that this function and the others could instead be static/class methods of their associated State class. Let's discuss it

clement-moulin-frier · 2024-03-28T14:43:04Z

vivarium/simulator/states.py

+        behavior: int = 1,
+        wheel_diameter: float = 2.,
+        speed_mul: float = 1.,
+        theta_mul: float = 1.,
+        prox_dist_max: float = 40.,
+        prox_cos_min: float = 0.,


Idem, we should be able to pass arrays if we want to

clement-moulin-frier · 2024-03-28T14:48:31Z

vivarium/simulator/states.py

+        theta_mul: float = 1.,
+        prox_dist_max: float = 40.,
+        prox_cos_min: float = 0.,
+        color: str = "blue"


To be discussed, I'm not sure it is super coherent to pass colors as strings, since we are in the simulator world here. For example, behavior above is defined as an int, not a string (whereas it is defined as a string in the client Configs, as for colors).
Also, specifying colors as a RGB vector here will solve the issue I mentioned about matlplotlib above. Maybe the best solution would be to have a RGB <-> String converter in a utils module, that can be easily imported wherever we need this. For example here ; or in the run script

clement-moulin-frier · 2024-03-28T14:49:20Z

vivarium/simulator/states.py

+    # TODO : Allow to define custom list of behaviors, wheel_diameters ... (in fact for all parameters)
+    # if the shape if just 1 value, assign it to all agents
+    # else, ensure you are given a list of arguments of size n_agents and transform it into a jax array


Great, we thought of the same thing :) I agree, let's do this

clement-moulin-frier · 2024-03-28T14:49:42Z

vivarium/simulator/states.py

+
+def init_object_state(
+        simulator_state: SimulatorState,
+        color: str = "red"


Same comment as above for color

clement-moulin-frier · 2024-03-28T14:54:42Z

vivarium/simulator/states.py

+def init_state(
+        simulator_state: SimulatorState,
+        agents_state: AgentState,
+        objects_state: ObjectState,
+        nve_state: NVEState
+        ) -> State:
+
+    return State(
+        simulator_state=simulator_state,
+        agent_state=agents_state,
+        object_state=objects_state,
+        nve_state=nve_state
+    )
+


A general comment to discuss:

If I understand well, with this PR a user will be able to define a custom state more easily from a run script, which is nice indeed. But if the objective is to simplify the life of the use, I am wondering if we shouldn't abstract the distinction between the NVEState and the Agent/ObjectState. How about to just having a function where a user can just enter what he wants (e.g. how many agents/objects, which positions etc ...) and then the function will deal with how to construct the State, hiding this complexity to the user.

…corentin/init_state_without_configs

corentinlger added 9 commits March 13, 2024 12:41

First refactoring step save

ec5c2a8

Merge branch 'main' of github.com:clement-moulin-frier/vivarium into …

4534c5d

…corentin/init_state_without_configs

Add all init functions in state.py

98046d4

Fix imports with new states file

b498f04

Initialize states without configs

89ecef7

Fix init state values that turn into NaN in client side

4308ce6

Fix converters import error

868074e

Change arguments for agents and objects states init

b86da7a

Add option to test initialization with custom positions

225c5f6

corentinlger requested a review from clement-moulin-frier March 20, 2024 16:01

corentinlger added 3 commits March 20, 2024 17:33

Update tests for new state initialization

4ee97dc

Merge branch 'main' of github.com:clement-moulin-frier/vivarium into …

d0e62e8

…corentin/init_state_without_configs

Reduce default value of agents proximeters range

b7e98d9

corentinlger requested a review from Marsolo1 March 21, 2024 09:40

corentinlger added 2 commits March 21, 2024 11:50

Refactor init_positions function and update comments

fd96746

Allow defining a initial number of existing entities

98a505b

Marsolo1 reviewed Mar 25, 2024

View reviewed changes

Marsolo1 approved these changes Mar 25, 2024

View reviewed changes

clement-moulin-frier reviewed Mar 28, 2024

View reviewed changes

corentinlger added 5 commits April 4, 2024 00:00

Remove code to test custom positions in run_server

0d42317

Merge branch 'main' of github.com:clement-moulin-frier/vivarium into …

5bc67f3

…corentin/init_state_without_configs

Fix vmap bug on normal function and clean sim_computation

adefcce

Add collision parameters and clean files

a54f920

Fix scipy version

efa959e

corentinlger merged commit a44e5ed into main Apr 4, 2024
2 checks passed

corentinlger mentioned this pull request Apr 4, 2024

Improve PR 57 #71

Closed

corentinlger deleted the corentin/init_state_without_configs branch April 16, 2024 15:45



		def test_init_simulator_args():
		""" Test the initialization of state with arguments """

		return jnp.array(list(mcolors.to_rgb(color_str)))


		def init_simulator_state(



		def _init_positions(key_pos, positions, n_elements, box_size, n_dims=2):
		assert (len(positions) == n_elements if positions else True)

		key_pos, key_or = random.split(key)
		key_ag, key_obj = random.split(key_pos)

		use_fori_loop=jnp.array([1*use_fori_loop]))


		def _init_positions(key_pos, positions, n_elements, box_size, n_dims=2):

Init state without configs #57

Init state without configs #57

Conversation

corentinlger commented Mar 20, 2024 • edited Loading

Description

Related Issue (if applicable)

How to Test

Screenshots (if applicable)

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Marsolo1 left a comment

Choose a reason for hiding this comment

clement-moulin-frier left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Marsolo1 commented Mar 28, 2024

clement-moulin-frier left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

corentinlger commented Mar 20, 2024 •

edited

Loading