code: def process_experience(
      experience: list[list[agent.ExpTuple]],
      actor_steps: int,
      num_agents: int,
      gamma: float,
      lambda_: float,
  ):
      """Process experience for training, including advantage estimation.
  
      Args:
        experience: collected from agents in the form of nested lists/namedtuple
        actor_steps: number of steps each agent has completed
        num_agents: number of agents that collected experience
        gamma: dicount parameter
        lambda_: GAE parameter
  
      Returns:
        trajectories: trajectories readily accessible for `train_step()` function
      """
      obs_shape = (84, 84, 4)
      exp_dims = (actor_steps, num_agents)
      values_dims = (actor_steps + 1, num_agents)
      states = np.zeros(exp_dims + obs_shape, dtype=np.float32)
      actions = np.zeros(exp_dims, dtype=np.int32)
      rewards = np.zeros(exp_dims, dtype=np.float32)
      values = np.zeros(values_dims, dtype=np.float32)
      log_probs = np.zeros(exp_dims, dtype=np.float32)
      dones = np.zeros(exp_dims, dtype=np.float32)
  
      for t in range(
          len(experience) - 1
      ):  # experience[-1] only for next_values
          for agent_id, exp_agent in enumerate(experience[t]):
              states[t, agent_id, ...] = exp_agent.state
              actions[t, agent_id] = exp_agent.action
              rewards[t, agent_id] = exp_agent.reward
              values[t, agent_id] = exp_agent.value
              log_probs[t, agent_id] = exp_agent.log_prob
              # Dones need to be 0 for terminal states.
              dones[t, agent_id] = float(not exp_agent.done)
      for a in range(num_agents):
          values[-1, a] = experience[-1][a].value
      advantages = gae_advantages(
          rewards, dones, values, gamma, lambda_
      )
      returns = advantages + values[:-1, :]
      # After preprocessing, concatenate data from all agents.
      trajectories = (states, actions, log_probs, returns, advantages)
      trajectory_len = num_agents * actor_steps
      trajectories = tuple(
          map(
              lambda x: np.reshape(x, (trajectory_len,) + x.shape[2:]),
              trajectories,
          )
      )
      return trajectories
  
imports: import functools
  from typing import Any, Callable
  from absl import logging
  import flax
  from flax import linen
  import agent
  import models
  import test_episodes
  from flax.metrics import tensorboard
  from flax.training import checkpoints
  from flax.training import train_state
  import jax
  import jax.numpy
  import ml_collections
  import numpy
  import optax
label_desc: Use image node_id 5 for calling a global function with example usage: process_experience(experience, actor_steps, num_agents, gamma, lambda_) and returns: trajectories