def policy_rollout(agent, env, num_traj, horizon):
  """Rollout an agent to collect trajectories.

  Args:
    agent: an agent to rollout.
    env: an environment to perform rollouts.
    num_traj: the number of trajectories to collect.
    horizon: the maximal number of steps for each trajectory.

  Returns:
    states, actions, rewards and observations from rollout trajectories.
  """
  traj_states = []
  traj_actions = []
  traj_rewards = []
  traj_observations = []

  for _ in range(num_traj):
    time_step = env.reset()
    states = []
    rewards = []
    actions = []
    observations = []

    for _ in range(horizon):
      # MuJoCo specific operations.
      states.append(env._gym_env.get_state())  # pylint: disable=protected-access
      observations.append(time_step)

      action = agent.action(time_step.observation)
      actions.append(action)
      time_step = env.step(action)
      rewards.append(float(time_step.reward))

      if time_step.is_last():
        break

    traj_states.append(states)
    traj_actions.append(actions)
    traj_rewards.append(rewards)
    traj_observations.append(observations)

  return traj_states, traj_actions, traj_rewards, traj_observations