Once the environments are vectorized, how do I have to gather immediate experiences for the agent?

Question

My main purpose right now is to train an agent using the A2C algorithm to solve the Atari Breakout game. So far I have succeeded to create that code with a single agent and environment. To break the correlation between samples (i.i.d), I need to have an agent interacting with several environments.

class GymEnvVec():

    def __init__(self, env_name, n_envs, seed=0):
        make_env = lambda: gym.make(env_name)
        self.envs = [make_env() for _ in range(n_envs)]
        [env.seed(seed + 10 * i) for i, env in enumerate(self.envs)]

    def reset(self):
        return [env.reset() for env in self.envs]

    def step(self, actions):
        return list(zip(*[env.step(a) for env, a in zip(self.envs, actions)]))

I can use the class GymEnvVec to vectorize my environment.

So I can set my environments with

envs = GymEnvVec(env_name="Breakout-v0", n_envs=50)

I can get my first observations with

observations = envs.reset()

Pick some actions with

actions = agent.choose_actions(observations)

The choose_actions method might look like

def choose_actions(self, states):
        assert isinstance(states, (list, tuple))

        actions = []
        for state in states:
            probabilities  = F.softmax(self.network(state)[0])
            action_probs = T.distributions.Categorical(probabilities)
            actions.append(action_probs.sample())

        return [action.item() for action in actions] 

Finally, the environments will spit the next_states, rewards and if it is done with

next_states, rewards, dones, _ = env.step(actions)

It is at this point I am a bit confused. I think I need to gather immediate experiences, batch altogether and forward it to the agent. My problem is probably with the "gather immediate experiences".

I propose a solution, but I am far from being sure it is a good answer. At each iteration, I think I must take a random number with

nb = random.randint(0, len(n_envs)-1)

and put the experience in history with

history.append(Experience(state=states[nb], actions[nb], rewards[nb], dones[nb]))

Am I wrong? Can you tell me what I should do?

Answer 1

class ExperienceSource():
    def __init__(self, env, agent, reward_steps):
        self.env = env
        self.agent = agent
        self.reward_steps = reward_steps

    def __iter__(self):
        histories = [deque(maxlen=self.reward_steps) for i in range(len(self.env.envs))]
        states = self.env.reset()

        while True:

            for idx, env in enumerate(self.env.envs):
                action = self.agent.choose_action(states[idx])
                state, reward, done, _ = env.step(action)

                current_rewards[idx] += reward
                histories[idx].append(Experience(state, action, reward, done))

                if len(histories[idx]) == self.reward_steps:
                    yield tuple(histories[idx])

                if done: 
                    yield tuple(histories[idx])
                    state = env.reset()

Be aware that self.reward_steps is simply the value defined by N-1 in the following formula $$Q(s,a) = \sum_{i=0}^{N-1} \gamma^i r_i + \gamma^N V(s_N)$$ and self.env is simply an instance of GymEnvVec class from the question.