pickling error when multiprocessing for batched continuation of environments

[MiladInk]

Hi All!

Thanks for the amazing melting pot environments. I am trying to train agents using jax on the environments. To enable this, I keep a batch of environments - which is just a list of let's say 32 environments - and I call the step function on all of them in a for loop. To make it a bit faster I wanted to call the step function for each one using multiprocessing. But I got cannot pickle 'dmlab2d.dmlab2d_pybind.Lab2d' object. I am wondering if you have suggestions for me for what I want to do.

[jagapiou]

No good suggestions sorry.

Looks like you'll have to find a way to avoid pickling Lab2d objects. Perhaps try just pickling the substrate name and have the processes spin up the environment themselves?

If you want Lab2d to support pickling you'll need to make a feature request at https://github.com/google-deepmind/lab2d. I expect they won't have time for this I'm afraid.

[MiladInk]

Oh, I see. I just cannot make the process spin up its own environment as I want agents to interact with a batch of environments back and forth. Thanks for the swift reply.

[jagapiou]

Haven't played with this myself, but I think maybe you can set up one process for each environment that constructs the Lab2d itself and is communicated with by the main process only via its env.step(actions). Then I think there will be no need to serialize the Lab2d (only the actions/timestep).