class_docstr: Manages the state of a single agent/environment interaction loop.
docstr: Run unroll_length agent/environment steps, returning the trajectory.
code: def unroll(
      self,
      rng_key,
      frame_count: int,
      params: hk.Params,
      unroll_length: int,
  ) -> util.Transition:
      """Run unroll_length agent/environment steps, returning the trajectory."""
      timestep = self._timestep
      agent_state = self._agent_state
      # Unroll one longer if trajectory is empty.
      num_interactions = unroll_length + int(not self._traj)
      subkeys = jax.random.split(rng_key, num_interactions)
      for i in range(num_interactions):
          timestep = util.preprocess_step(timestep)
          agent_out, next_state = self._agent.step(
              subkeys[i], params, timestep, agent_state
          )
          transition = util.Transition(
              timestep=timestep,
              agent_out=agent_out,
              agent_state=agent_state,
          )
          self._traj.append(transition)
          agent_state = next_state
          timestep = self._env.step(agent_out.action)
  
          if timestep.last():
              self._episode_return += timestep.reward
              self._logger.write(
                  {
                      "num_frames": frame_count,
                      "episode_return": self._episode_return,
                  }
              )
              self._episode_return = 0.0
          else:
              self._episode_return += timestep.reward or 0.0
  
          # Elide a manual agent_state reset on step_type.first(), as the ResetCore
          # already takes care of this for us.
  
      # Pack the trajectory and reset parent state.
      trajectory = jax.device_get(self._traj)
      trajectory = jax.tree_util.tree_map(
          lambda *xs: np.stack(xs), *trajectory
      )
      self._timestep = timestep
      self._agent_state = agent_state
      # Keep the bootstrap timestep for next trajectory.
      self._traj = self._traj[-1:]
      return trajectory
  
file_path: dm-haiku/examples/impala/actor.py
imports: import dm_env
  import haiku
  from examples.impala import agent
  from examples.impala import learner
  from examples.impala import util
  import jax
  import numpy
label_desc: Use image node_id 2 for calling the Actor obj's underlying member method code with example usage: obj.unroll(rng_key, frame_count, params, unroll_length) and returns: trajectory