Avoiding static arguments when using JAX for MARL

[cool-RR]

Hi folks!

I'm using JAX for MARL experiments using an opponent shaping algorithm called POLA. This algorithm is complicated and takes a long time to compile, between minutes and 12 hours, depending on options. I believe that one of the causes for how long the compilations are taking is that I have many jitted functions that get an agent index as an argument, and that argument is marked as static.

Here is a grossly simplified example of code that I hope helps explain the issue: https://gist.github.com/cool-RR/dd0b548c1990ec111e03c5b8e0863808

The Population class has multiple agents, i.e. multiple TrainState objects in a list. The update_agent method trains one agent while depending on the behavior of all the other agents. As you can see on line 38, the agent argument is marked as static. If I don't mark it as static, JAX would raise an exception, since I'm using agent as an array index later.

This is a simplified example; in my real code I have a dozen different methods that call each other and they all have an agent argument that's marked as static.

Am I right in my intution that this static argument is making the compilation of my program slower? Is there any way for me to make it so the agent argument would not be static?

Thanks for your help,
Ram Rachum.

[jakevdp]

As written, you will have to re-compile your update_agent for each unique value of the integer agent (JIT compilation caches are keyed based on the static arguments). This will likely cause a significant slowdown if you're working with a large number of agents.

A better pattern would be to avoid having a list of train states, but instead have a single batched train state which you can index into using a dynamic agent index. In other words, use a struct-of-arrays pattern for your train states rather than the current list-of-structs pattern. A further benefit of this would be that you could replace the for loop in loss_fun with a vmap followed by some post-processing of the outputs, which would speed up the compilation and runtime even further.

Does that make sense?

[jakevdp]

Sorry, I must have misunderstood what you meant when you mentioned pytree_ravel (I assume what you're referring to is jax.flatten_util.ravel_pytree)? I'll take a look at your code, but I'm not sure I'll have the bandwidth to dig into diffs between your branches

Tree map indexing is not as slow as ravelling and unravelling: it's just applying indexing to the batched arrays within the training data structure.

[jakevdp]

Sorry, I opened your repository, and I have no idea where I'm supposed to be looking. There are dozens of files there.

[jakevdp]

General compilation time debugging tips: any time you're using Python control flow (e.g. a for loop, a while loop, or recursive function calls) you're likely generating many XLA operations in sequence, because JAX flattens all Python control flow. The more XLA operations you have, the longer your program will be, and the longer your program is, the longer it will take to compile (compile time is generally super-linear in the number of operations)

So to fix slow compilation, first look for places where you have for loops within a jit compiled function, and either replace them with vmap-based batching if the operations have no dependence between batching, or with fori_loop/scan in cases where operations have sequential dependence.

This is why my first recommendation to you was to avoid the list-of-structs (which require looping via Python control flow) in favor of struct-of-arrays (which allow batching via vmap).

[jakevdp]

Yes, it's pretty big. pola_algorithm.py is probably the most important file. The Population.train method there is kind of the main function. The Population class used to contain a field brains which was a tuple of Brain objects. (A Brain object has a pair of train states, one for actor and one for critic, and a bunch of metadata.) After the change, the Population class has an attribute raveled_brain_by_agent which is an array, and then Population.__getitem__ returns self.brain_unraveler(self.raveled_brain_by_agent[i]) so you could get brains by doing population[agent].

Yeah - it looks like you have thousands of lines of existing code, and it seems like that may be tricky to refactor the fundamental data structures used in that code. But the general advice above will apply: avoid Python control flow if you find that compilation times are getting out of hand.

[cool-RR]

I understand, thank you. I'll think on how to proceed.

[cool-RR]

Yes, it's pretty big. pola_algorithm.py is probably the most important file. The Population.train method there is kind of the main function. The Population class used to contain a field brains which was a tuple of Brain objects. (A Brain object has a pair of train states, one for actor and one for critic, and a bunch of metadata.) After the change, the Population class has an attribute raveled_brain_by_agent which is an array, and then Population.__getitem__ returns self.brain_unraveler(self.raveled_brain_by_agent[i]) so you could get brains by doing population[agent].