code: def get_agent_timestep(
      timestep: dm_env.TimeStep, agent_id: types.AgentID
  ) -> dm_env.TimeStep:
      """Returns the extracted timestep for a particular agent."""
      # Discounts are assumed to be shared amongst agents
      reward = (
          None if timestep.reward is None else timestep.reward[agent_id]
      )
      return dm_env.TimeStep(
          observation=timestep.observation[agent_id],
          reward=reward,
          discount=timestep.discount,
          step_type=timestep.step_type,
      )
  
imports: from acme import specs
  from acme.multiagent import types
  import dm_env
label_desc: Use image node_id 2 for calling a global function with example usage: get_agent_timestep(timestep, agent_id) and returns: dm_env