Pmaping RL experiments

[joaogui1]

So, I'm trying to use deluca to parallelize my Reinforcement Learning loop, the idea is to run the same agent-environment pair with different random seeds so I can get mean and standard deviation of rewards. Right now I'm creating a list of agents, each with a different random seed, and trying to broadcast the environment, but when I use pmap I get the following error:

ValueError: pmap got arg 0 of rank 0 but axis to be mapped 0. The tree of ranks is:
((<deluca.envs.classic._cartpole.CartPole object at 0x7f9bc10fa410>, <deluca.agents._deep.Deep object at 0x7f9bc0c87810>), {})

Here's a link to the colab reproducing the error.
Do I have to reshape things? Doesn't pmap work with lists? Or is the problem that the agents don't count as pytrees? Can I register then as one?

[davisyoshida]

I don't know about the overall question, but regarding using them as pytree nodes, there's a register_pytree_node function. The documentation for the tree_util package seems to be missing in action though (https://jax.readthedocs.io/en/latest/jax.tree_util.html).