Baselines: Where does the aggregation happen?

[richardrl]

When running multiple workers, where in the code does aggregating the rollouts/gradients happen?

[richardrl]

I expect somewhere there is a driver directing all the threads, and in that driver it aggregates (sums) the gradients and then does backprop on a shared model between the threads... where is that logic happening?

[pzhokhov]

HI @richardrl! That would be SubprocVecEnv class (or ShmemVecEnv), located here: https://github.com/openai/baselines/blob/5b41c926c7a852df3f0928afdf2429f96a3965cb/baselines/common/vec_env/subproc_vec_env.py#L34
and here: https://github.com/openai/baselines/blob/5b41c926c7a852df3f0928afdf2429f96a3965cb/baselines/common/vec_env/shmem_vec_env.py#L20. Note that these aggregate rollouts (not gradients) from workers (in other words, each worker subprocess performs only environment steps; nothing else)
On the other hand, aggregation of gradients from multiple workers is possible when running with multiple MPI processes. In that case,
aggregation logic is here: https://github.com/openai/baselines/blob/5b41c926c7a852df3f0928afdf2429f96a3965cb/baselines/common/mpi_adam_optimizer.py#L5

[richardrl]

Thanks @pzhokhov

I'm trying to follow the logic with the HER example.

So starting from here:
mpirun -np 19 python -m baselines.run --num_env=2 --alg=her ...

We eventually end up with each worker running the train function in her.py. On line 41, we have this logic:

for _ in range(n_cycles):

            episode = rollout_worker.generate_rollouts()

            policy.store_episode(episode)

            for _ in range(n_batches):

                policy.train()

the 'policy' variable here is an instance of DDPG.

In DDPG, the train() function eventually calls self._grads(). I would assume here is where aggregation happens but there is no indication in self._grads() that any distributed aggregation over processor threads is happening?

def _grads(self):
        # Avoid feed_dict here for performance!
        critic_loss, actor_loss, Q_grad, pi_grad = self.sess.run([
            self.Q_loss_tf,
            self.main.Q_pi_tf,
            self.Q_grad_tf,
            self.pi_grad_tf
        ])
        return critic_loss, actor_loss, Q_grad, pi_grad

[richardrl]

Also, does the rollout worker get its policy updated from a master policy after every rollout? I can't seem to find that logic in the code

[pzhokhov]

ah I see. In that case the logic is in MpiAdam (MpiAdam is used here https://github.com/openai/baselines/blob/5b41c926c7a852df3f0928afdf2429f96a3965cb/baselines/her/ddpg.py#L389 and next line), and its internal logic reponsible for both gradient aggregation and subsequent updates of worker parameters is here (https://github.com/openai/baselines/blob/5b41c926c7a852df3f0928afdf2429f96a3965cb/baselines/common/mpi_adam.py#L10)

[richardrl]

Hi @pzhokhov

To clarify:
In the HER "mpirun np -19 ... --num_env 2" example for reproducing the paper results...

SubprocVecEnv allows a single thread to rollout from 2 environments simultaneously. These rollouts are collected in a replay buffer that is local to this thread.

Gradients are computed locally for each thread, from each thread's local replay buffer. Then the gradients are aggregated, added to the current policy, and the new policy is redistributed to each thread. In other words, the only inter-thread communication is the gradient updates; the replay buffers are separate objects for each thread, and the gradients computed per thread are always based on samples from that thread's local replay buffer. Is that correct?