code: def test_unroll(self):
      mock_learner = mock.MagicMock()
      traj_len = 10
      actor = actor_lib.Actor(
          agent=self.agent,
          env=self.env,
          learner=mock_learner,
          unroll_length=traj_len,
      )
      self.key, subkey = jax.random.split(self.key)
      act_out = actor.unroll(
          rng_key=subkey,
          frame_count=0,
          params=self.initial_params,
          unroll_length=traj_len,
      )
  
      self.assertIsInstance(act_out, util.Transition)
      self.assertIsInstance(act_out.timestep, dm_env.TimeStep)
      self.assertLen(act_out.timestep.reward.shape, 1)
      self.assertEqual(act_out.timestep.reward.shape, (traj_len + 1,))
      self.assertLen(act_out.timestep.discount.shape, 1)
      self.assertEqual(act_out.timestep.discount.shape, (traj_len + 1,))
      self.assertLen(act_out.timestep.step_type.shape, 1)
      self.assertEqual(
          act_out.timestep.step_type.shape, (traj_len + 1,)
      )
  
      self.assertLen(act_out.timestep.observation.shape, 3)
      self.assertEqual(
          act_out.timestep.observation.shape,
          (traj_len + 1,) + self.obs_spec.shape,
      )
  
      self.assertIsInstance(act_out.agent_out, agent_lib.AgentOutput)
      self.assertLen(act_out.agent_out.action.shape, 1)
      self.assertEqual(act_out.agent_out.action.shape, (traj_len + 1,))
  
      self.assertLen(act_out.agent_out.policy_logits.shape, 2)
      self.assertEqual(
          act_out.agent_out.policy_logits.shape,
          (traj_len + 1, self.num_actions),
      )
  
      self.assertLen(act_out.agent_out.values.shape, 1)
      self.assertEqual(act_out.agent_out.values.shape, (traj_len + 1,))
  
      self.assertEqual(act_out.agent_state.shape, (traj_len + 1,))
  
file_path: dm-haiku/examples/impala/actor_test.py
imports: from unittest import mock
  from absl.testing import absltest
  from bsuite.environments import catch
  import dm_env
  from examples.impala import actor
  from examples.impala import agent
  from examples.impala import haiku_nets
  from examples.impala import learner
  from examples.impala import util
  import jax
  import numpy
  import tree
label_desc: Use image node_id 2 for calling the CatchTest obj's underlying member method code with example usage: obj.test_unroll() without return types