Runs a batch of rollouts from given initial observations.

def run_rollouts(
    env, agent, initial_observations, step_limit=None, discount_factor=1.0,
    log_every_steps=None, video_writers=(), color_bar=False,
    many_rollouts_from_each_env=False
):
  """Runs a batch of rollouts from given initial observations."""
  assert step_limit is not None or not many_rollouts_from_each_env, (
      "When collecting many rollouts from each environment, time limit must "
      "be set."
  )

  num_dones = 0
  first_dones = np.array([False] * env.batch_size)
  observations = initial_observations
  step_index = 0
  cum_rewards = np.zeros(env.batch_size)

  for (video_writer, obs_stack) in zip(video_writers, initial_observations):
    for (i, ob) in enumerate(obs_stack):
      debug_frame = augment_observation(
          ob, reward=0, cum_reward=0, frame_index=(-len(obs_stack) + i + 1),
          bar_color=((0, 255, 0) if color_bar else None)
      )
      video_writer.write(debug_frame)

  def proceed():
    if step_index < step_limit:
      return num_dones < env.batch_size or many_rollouts_from_each_env
    else:
      return False

  while proceed():
    act_kwargs = {}
    if agent.needs_env_state:
      act_kwargs["env_state"] = env.state
    actions = agent.act(observations, **act_kwargs)
    (observations, rewards, dones) = env.step(actions)
    observations = list(observations)
    now_done_indices = []
    for (i, done) in enumerate(dones):
      if done and (not first_dones[i] or many_rollouts_from_each_env):
        now_done_indices.append(i)
        first_dones[i] = True
        num_dones += 1
    if now_done_indices:
      # Unless many_rollouts_from_each_env, reset only envs done the first time
      # in this timestep to ensure that we collect exactly 1 rollout from each
      # env.
      reset_observations = env.reset(now_done_indices)
      for (i, observation) in zip(now_done_indices, reset_observations):
        observations[i] = observation
    observations = np.array(observations)
    cum_rewards[~first_dones] = (
        cum_rewards[~first_dones] * discount_factor + rewards[~first_dones]
    )
    step_index += 1

    for (video_writer, obs_stack, reward, cum_reward, done) in zip(
        video_writers, observations, rewards, cum_rewards, first_dones
    ):
      if done:
        continue
      ob = obs_stack[-1]
      debug_frame = augment_observation(
          ob, reward=reward, cum_reward=cum_reward,
          frame_index=step_index, bar_color=((255, 0, 0) if color_bar else None)
      )
      video_writer.write(debug_frame)

    # TODO(afrozm): Clean this up with tf.logging.log_every_n
    if log_every_steps is not None and step_index % log_every_steps == 0:
      tf.logging.info("Step %d, mean_score: %f", step_index, cum_rewards.mean())

  return (observations, cum_rewards)