Creating multiple workers to do DQN

[raymondchua]

Hi all, I hope this is the right place to ask this question. I would like be able to collect experience from 64 parallel environment copies into a single replay buffer. For exampling, using 8 parallel workers, using shared model parameters and averaging gradients between workers. Does anyone know if this can be done in Jax using vmap and jit? Perhaps similar to this: #11565 (reply in thread)

[anh-tong]

Hi,

It seems like you can pmap in this case. In my own experience, I use some of Flax utility functions (I think you can write your own similar function with no difficulty).

For examples,

• flax.jax_utils.replicate to make shared model parameters available to all devices
• flax.training.common_utils.shard to share data to all devices
• jax.lax.pmean: use to average computed gradients

Here is a simple structure I use

def compute_fn(params, input):
    grad = <compute grad here>
    
    # note that `axis_name` should be the same with `pmap`
    grad = jax.lax.pmean(grad, axis_name="batch")
    
    return grad

# make model available to all devices
params = <model>
params = replicate(params)

# make `pmap` version of `compute_fn`
p_compute_fn = jax.pmap(compute_fn, axis_name="batch", donate_argnums=(0,))

# shard data across all devices
data = <input data>
data = shard(data)

grad = p_compute_fn(params, data)

# `grad` is still replicated across the devices, can be retrieved as
grad = unreplicate(grad)

This is what I usually do. Maybe there is a better approach out there.

Hope this could help.