RatInABox2.0 - Roadmap

[TomGeorge1234]

I've begun to think about 2.0. The reason is that there are are certainly a couple of choices I made early on in development which weren't optimal. Now could be a good time to fix these as the community is growing but still small enough it won't be super disruptive. Also fixing them will make it easier to maintain RiaB in the long run.

I'm opening this issue to get community thoughts on this. @SynapticSage @colleenjg @jquinnlee @mehulrastogi you're some of the most active users I know fairly well so I'm tagging you to get your input (if you have any), but anyone can chip in here. Here's my thoughts:

Essential and backwards incompatible changes (do first):

• Refactoring: As discussed in major refactoring #58 Package structuring #55. E.g. it's not nice having all Neurons classes in one .py file.
• Args not Dicts: It's increasingly annoying me that parameters are always handed in as dicts. This is unconventional and has warranted very-well-made but hacky work-arounds e.g. Unnoticed, incorrectly typed out parameter names #38 Added params checking for all classes #39
• Global Environment update(): Given, now, Environments know about their Agents and Agents know about their Neurons we could have just one update function in Env which cascades through else thing else. Cleaner?
• Rename dev --> main
• Environment stores the global clock. This just makes sense imo.
• Better policy API - I don't love the drift_velocity kwarg. Maybe instead Agents can have a policy() method which returns a drift - this would default to the random motion policy, unifying that too. Just something to consider.
• get_state() repeats logic each time. Seems better to have a unique .forward() (or maybe called calculate_firing_rate()) method for each class which receives arrays of positions and head directions etc and a shared .get_state() which lives in Neurons which calls .forward(). I'm returning to +1 this idea, it makes a lot more sense. Also, instead of get_state() we can have numerous such as get_agent_firing_rate(), get_rate_map(), get_angular_tuning_curve() etc. These "get"-functions should return not just lists of the firing rates but also the lists of the respective coordinates (we could maybe use xarray for this but I don't like the extra dependency).
• EnvironmentEntity class could be made which can be added to an Environment which will then update each entity at each time step. Entities will have their own .render methods so Env can loop over them and plot them too. This could be used to flexibly build Teleports, Doors, Keys, etc. etc. for more dynamic environments.

Other essential changes

• Type hinting: This is a new thing in python which I've been told to consider. Any thoughts?
• Modularity Break down some of the larger update() perhaps adding into new agent/neuron/env specific utils scripts.
• Dynamic environments Environments can change by adding walls and objects but we should formalise this with setters which, whenever called, save the "state" of the environment alongside a timestamp as a dictionary to an Env.history dictionary. Then, when plotting / animating the environment we can pass in a time argument and the correct state can be retrieved and plotted. The state of the environment only appends to history whenever it changes (e.g. a setter is called).
  • Related to the above, if an Environment changes half way through a simulation then animations will not support this since they always replot the last environment which is both wasteful and possibly wrong. To get around this whenever you call plot_environment() it can be passed a fig an ax and a new object which is a list/dict of plot objects, R which are all matplotlib.Artists already existing on the figure. The environment can store an equivalent list of plot objects and whenever this changes (e.g. a wall is added or an object is moved etc.) this change is logged then plotting can (i) get the list of plot objects corresponding to the correct time and (ii) compared it to the passed list, if they aren't equal then repot the env, otherwise don't bother. Something like that.
  • Alternatively (maybe better): Environments have an Env.history dictionary storing the full "state" of the environment (all object locations, walls, boundaries, etc.). Then Env.plot_environment() takes a time argument and find the state of the at that time and plots that.
• Optimise Agent.update() as shown in the paper this is a significant bottleneck.
• Plotting: To me at least the visualisation ability of RiaB is really important but animations are slow and I like animating things so this annoying. Could improve by being smarter about how we render stuff in matplotlib, and not re-rendering the Environment or trajectories each frame. Stuff like that. See Slow animations #54
  • For example if this Environment state dictionary was stored inside the figure itself with some kind of hash code we could just check on each call whether the desired state matches the state thats been plotted. Only replot if they aren't equal.
• Align animate_ API with plot_ API just with a few extra kwargs.
• Only pass ax not fig to figure plotting functions. This may throw up some things but likely minor.
• Break up utils.py into separate ones for the Agent package, Neurons package and Env package and maybe also a misc.
• Documentation: Would be great to have a sustainable ReadTheDocs page. We should think about how to structure doc strings so they are all uniform. Documentation #36
• Unit testing: I have been pretty sloppy about this but will add loads more.
• Testing on PRs Run RiaB tests, test doc strings and text styling.
• Move this to RatInABox/RatInABox not TomGeorge1234/RatInABox
  • RatInABox/RatInABox_RL** package containing all the RL stuff (Actor, Critic, ValueNeuron, TDError, TaskEnv etc.)
• IntermediateNeurons subclass for neurons which aren't "fundamental" but take other neurons as inputs. Current examples are FeedForwardLayer and NeuralNetworkNeurons
• DynamicNeurons subclass for neurons which aren't static i.e. you can't call Neurons.plot_rate_map() because they actually depend on the past history. Examples include TDErrorNeurons (to be made) or anything with recurrency.
• SmoothRandomFeatureNeurons just some spatially tuned but random neurons. Users just provide a length scale. Would be useful for a lot of feature learning studies. Probably something like a gaussian process underlying these neurons.
• Social behaviour demos social behaviour seems like a low hanging fruit for neuroscience and RiaB
• Improve integration of RiaB into other toolkits we should actively look to contributing or collaborating with other toolkits (NeuroML, movement, DLC, OpenSourceNeuro etc.) to embed RiaB as in house motion lightweight data simulators.
• Fewer args more kwargs Some of the functions (in particular plotting) have quite bloated argument lists. I think it would be better to remove some of these and allow them to be hidden in **kwargs then defined at the top of the function arg = kwargs.get("arg_name",default_val). This is backwards compatible, cleans up the doc strings so more readable, and we can use this to expose any/all free parameters (even ones which weren't before hand an argument) for greater flexibility. I have done this to the Agent.plot_trajectory()
• Integer spikes Currently spikes are sampled as 0 or 1 per time bin. This is restrictive and means that high sampler rates are needed to resolve high firing rates. Could get around this by doing proper Poisson sampling and allowing n_spikes > 1. See Sampling multiple spikes in large DT regime? #131 for discussion.
• Proper rate map visualisation with KDE rather than using np.histogram2d. See issue Improving how RatInABox estimates and facilitates rate map estimation from observed data #125

Things to consider

• Neurons should follow torch.nn.module API - this would make more efficient the evaluation of complex feedforward graphs which currently happens in a backwards manner. This might require renaming the .get_state() method with .forward(). Need to think more about this
• conda Once all of the above is done it would be nice to publish this on the condo-forge channel.
• Jax compatibility: Very on the fence about this one. Probably leaning towards not doing it. Would be great to have speed ups, autograd and gpu capacity but it could be just a bit too much / unnecessary / off-putting for non-python geeks (tbh, like me). But if jax is the future I want to consider it. Options include:
  • Don't do it
  • Partial jax to hit a few heavy-lifting utils functions. Q: Does this even work, would converting to/from jax arrays not be inconveniently slow here?
  • Full jax no numpy. np--> jnp everywhere.
  • Both jax and numpy. Users choose which backend. This should hard but I've played around and probably could be done. Has complications though.
• Speed distribution could be changed from Rayleigh to Gamma so it support variance (see straight trajectory with fixed speed mean, but Ag.history['vel'] is not constant #113)

I'm not a software guy so @SynapticSage @mehulrastogi feel free to give high level comments about best way to go forward. A

[SynapticSage]

args, not dicts 👍

A helpful case study in support of args ...

Most of you (I'm sure) have seen Grant Sanderson's beautiful 3blue1brown YouTube channel. Grant impressively homebrewed the manim package that creates his stunning math videos.

Similar to here, manim started out using a CONFIG dict. On a positive note, the CONFIG dict cut down on lines in object init; encouraged people to spell out settings in one place. But on the dark side, dicts required nearly re-coding a lot of Python features handled by kwargs and setting attributes. Ultimately the community fork decided to kill CONFIG dicts in favour of args -- decision convo here: https://github.com/ManimCommunity/manim/pull/763

Grant Sanderson's fork is also trying to remove them: 3b1b/manim#1932

plotting 👍

💯 replotting = slow.

... if ratinabox caches plot objects, super recommend scheme we chatted about:
#30 (comment)

The TaskEnvironment has a weak version of this feature -- doesn't replot everything and thus renders quickly. But it's pretty hacky in my view that the environment caches things about its agents and goals. In the long-run, it will be more maintainable to have each class in charge of caching its own plot objects rather than having to change master supervisor class's plot every time the children classes change.

type hinting 👍

Especially easy-to-type variables.

Tools like jedi and language-server-protocol offer better code completion for type-hinted variables.

unit testing 👍

global environment 👍

Possible suggestion: each RIB class could have a list of children (environment.childen = [agent, ...]; agent.children=[neuron,...]) to unify the way .update() and .plot()/.render() calls cascade down a hierarchy. It may be more uniform than each object having a different attribute name for its children.

Jax 🤷‍♂️

No strong opinions. Leaning partial Jax if the penalty for binary-op/shuttling numpy to a CPU jax.device is low.

[colleenjg]

Sounds like a great idea overall for the longevity of the package! I definitely agree for the args instead of dicts, type hinting and unit testing. For global environment, if the cascading update is implemented, I would suggest having a kwarg like cascade=True, to allow users to opt out, when needed. No strong views on the other sections.

I would suggest an additional section:
modularity: Many of the classes have very long methods that chain a lot of complex, separate computations together. When I've created new classes for my own use, e.g., new Agent classes, I've had to copy long sections of certain methods that I needed to overwrite, but only partially (for example, for computing an agent's velocity). This can create a lot of code duplication (I think there may already be some for the plotting methods). So, I strongly recommend adding the goal of modularization to the list, i.e., extracting meaningful subparts of class methods and turning them into their own functions, perhaps aggregated into agent_util.py, env_util.py and neuron_util.py, or something like that.

[TomGeorge1234]

Update: By the way I have now fully modularised Agent.update() as you suggested. I agree it's much more readable and flexible. Happy to use this as a model for future modularisation, thanks for the suggestion

[TomGeorge1234]

Great comments, thanks guys. @SynapticSage 3B1B advice heeded! @colleenjg you're right this could be more modular, for example Agent.update() is pretty enormous. Breaking these down would make sense so I'll look to do that. Don't expect this anytime soon btw so any new ideas, keep posting them here.

[jquinnlee]

These all sound like great changes for RAIB 2.0, and I agree w all of the comments from @SynapticSage and @colleenjg :)

I'm a particularly big fan of the global environment updating, as this seems much more concise. My only concern is whether this would slow down updates for really long simulations (like the ones I have been running, e.g. @ 30 Hz x 31 sessions x 40 min/session). It might be ideal to perform more selective updates and skip others if they are going to be static using some sort of argument in update()?

As far as Jax compatibility, I would be very much for this if it can actually speed things up for the heavier computations and long simulations, but as you point out it might not save compute time if large arrays are being converted often. I believe it would be worth some case testing in a couple of large simulations before ruling this out.

[TomGeorge1234]

Thanks for the feedback, closing for now.

[colleenjg]

One thing that just occurred to me, which could be considered:

Only passing ax to the plotting functions, not fig.

In typical use cases, to my knowledge, passing both should be redundant, as you can access the figure with ax.figure (or ax.ravel()[0].figure in cases where ax is an array).

[TomGeorge1234]

Agreed and added to the list. It's essentially redundant and only add bloat

[musicinmybrain]

If you add jax support, would it be possible to do it through an optional extra for opportunistic speed-ups rather than as a mandatory dependency?

As the primary maintainer of the Fedora Linux package for this project, I’m not sure if packaging https://github.com/google/jax would be feasible for us or not. While it does look like jax can be built without support for the proprietary CUDA SDK, it’s still a pretty gnarly stack when taken together with https://github.com/openxla/xla, and I’m not sure whether or not an attempt to package it would end up hitting a hard requirement on something nonfree.

[TomGeorge1234]

@musicinmybrain thanks for your feedback - that's ok, I doubt we'd go full jax. In fact leaning towards no jax at all actually. After some preliminary testing seems like getting significant speed ups would be difficult as most of the heavy computations are already vectorised

[mehulrastogi]

We should consider type hinting for Riab 2.0. (https://mypy.readthedocs.io/en/stable/cheat_sheet_py3.html)

This will enable users to do easier lookups and help with the autocompletion of code for anyone coding in an integrated development environment. Will mostly involve 2 things:-

• changing the variable names for objects (function parameters and local vars) from camel case to lower cases (Agents->agents).
  • Should also take note of any dirty variable naming (for eg. super in OVCs Vector cell refactoring #77 now solved :) ) if they still exist
• Type hinting the variables

[TomGeorge1234]

@mehulrastogi just thinking...do you think a more scalable way to support dynamic environments would be for all environmental objects to be their own class. I can imagine a ratinabox.Walls object which stores the walls and has an update() function just like Agent() which stores current wall locations into its own history dataframe. Then the Env.plot_environment(t=None) just receives a time stamp and gets the wall locations from their history at that time stamp and plots them. Same for objects. Same for holes. Same for any new "thing" (e.g. Teleporter, Wormholes, RewardSpouts I don't know!) we want to add down the line.

This is as opposed to the solution I'd previously imagined where the environment itself store a "state" dictionary which can be updated. The nice thing about this new proposal is that Walls, Objects, Wormholes etc. can be individually responsible for their rendering function. The Environment just contains a list of things it knows it has to render and calls thing.render(t=t, fig=fig, ax=ax) on each of these things. Seems more OOPy?

My main concern is that it would create memory issues storing so much more data but to be honest I'm not sure it would be much more data than Neurons already save without problems. We could also image some memory-clever solution since 99% of the time walls won't move.

[mehulrastogi]

This is a great idea but there area few things to consider for this:-

• we would be making a rather simple static environment also hard to read and modify (which some users might not like). I think we. can go more OOPy and have a different environment altogether. One dynamic and one static. As I think both will be initialized in a different way. We would also need to provide a friendly API to make these (I was thinking of a playground of sorts to create the environments with GUI. although this is a very long-term plan. although it will be a good addition to dynamic env. imagine you being able to create mazes and place different objects and 'things' to interact with in a GUI which takes care of the environment code. and being able to just in the playground design paradigms including different stages and any changes to the environment without extensive coding)
• I love the idea of different things in the env that the agent can interact with. I am not that concerned with memory to be honest but what concerns me is the rendering. These things could slow down rendering by a lot.

[TomGeorge1234]

Good points. Thanks for the feedback