code: def test_cut(self):
      # Create an environment.
      episode_1, _ = self._mock_multi_agent_records_from_env(size=100)
  
      # Assert that the episode has 100 timesteps.
      self.assertEqual(episode_1.t, 100)
  
      # Create a successor.
      episode_2 = episode_1.cut()
      # Assert that it has the same id.
      self.assertEqual(episode_1.id_, episode_2.id_)
      # Assert that all `SingleAgentEpisode`s have identical ids.
      for agent_id, agent_eps in episode_1.agent_episodes.items():
          self.assertEqual(
              agent_eps.id_, episode_2.agent_episodes[agent_id].id_
          )
      # Assert that the timestep starts at the end of the last episode.
      self.assertEqual(episode_1.t, episode_2.t_started)
      # Assert that the last observation and info of `episode_1` are the first
      # observation and info of `episode_2`.
      for agent_id, agent_eps in episode_1.agent_episodes.items():
          # If agents are done only ensure that the `SingleAgentEpisode` is empty.
          if agent_eps.is_done:
              self.assertEqual(
                  len(episode_2.agent_episodes[agent_id].observations),
                  0,
              )
          else:
              self.assertEqual(
                  agent_eps.observations[-1],
                  episode_2.agent_episodes[agent_id].observations[0],
              )
              self.assertEqual(
                  agent_eps.infos[-1],
                  episode_2.agent_episodes[agent_id].infos[0],
              )
  
      # Now test the buffers.
      for agent_id, agent_buffer in episode_1.agent_buffers.items():
          # Make sure the action buffers are either both full or both empty.
          self.assertEqual(
              agent_buffer["actions"].full(),
              episode_2.agent_buffers[agent_id]["actions"].full(),
          )
          # If the action buffers are full they should share the same value.
          if agent_buffer["actions"].full():
              self.assertEqual(
                  agent_buffer["actions"].queue[0],
                  episode_2.agent_buffers[agent_id]["actions"].queue[0],
              )
          # If the agent is not done, the buffers should be equal in value.
          if not episode_1.agent_episodes[agent_id].is_done:
              # The other buffers have default values, if the agent is not done.
              # Note, reward buffers could be full of partial rewards.
              self.assertEqual(
                  agent_buffer["rewards"].queue[0],
                  episode_2.agent_buffers[agent_id]["rewards"].queue[0],
              )
              # Here we want to know, if they are both different from `None`.
              self.assertEqual(
                  agent_buffer["extra_model_outputs"].queue[0],
                  episode_2.agent_buffers[agent_id][
                      "extra_model_outputs"
                  ].queue[0],
              )
          # If an agent is done the buffers should be empty for both, predecessor
          # and successor.
          else:
              self.assertTrue(agent_buffer["actions"].empty())
              self.assertTrue(agent_buffer["rewards"].empty())
              self.assertTrue(
                  agent_buffer["extra_model_outputs"].empty()
              )
              self.assertTrue(agent_buffer["actions"].empty())
              self.assertTrue(agent_buffer["rewards"].empty())
              self.assertTrue(
                  agent_buffer["extra_model_outputs"].empty()
              )
  
      # Ensure that the timestep mappings match.
      for (
          agent_id,
          agent_global_ts,
      ) in episode_2.global_t_to_local_t.items():
          # If an agent is not done, we write over the timestep from its last
          # observation.
          if not episode_2.agent_episodes[agent_id].is_done:
              self.assertEqual(
                  agent_global_ts[0],
                  episode_1.global_t_to_local_t[agent_id][-1],
              )
          # In the other case this mapping should be empty.
          else:
              self.assertEqual(len(agent_global_ts), 0)
  
      # Assert that the global action timestep mappings match.
      for (
          agent_id,
          agent_global_ts,
      ) in episode_2.global_actions_t.items():
          # If an agent is not done, we write over the timestep from its last
          # action.
          if not episode_2.agent_episodes[agent_id].is_done:
              # If a timestep mapping for actions was copied over the last timestep
              # of the üredecessor and the first of the successor must match.
              if agent_global_ts:
                  self.assertEqual(
                      agent_global_ts[0],
                      episode_1.global_actions_t[agent_id][-1],
                  )
              # If no action timestep mapping was copied over the last action must
              # have been at or before the last observation in the predecessor.
              else:
                  self.assertGreaterEqual(
                      episode_1.global_t_to_local_t[agent_id][-1],
                      episode_1.global_actions_t[agent_id][-1],
                  )
          # In the other case this mapping should be empty.
          else:
              self.assertEqual(len(agent_global_ts), 0)
  
      # Assert that the partial reward mappings and histories match.
      for (
          agent_id,
          agent_global_ts,
      ) in episode_2.partial_rewards_t.items():
          # Ensure that timestep mapping and history have the same length.
          self.assertEqual(
              len(agent_global_ts),
              len(episode_2.partial_rewards[agent_id]),
          )
          # If an agent is not done, we write over the timestep from its last
          # partial rewards.
          if not episode_2.agent_episodes[agent_id].is_done:
              # If there are partial rewards after the last observation ensure
              # they are correct.
              if (
                  episode_1.global_t_to_local_t[agent_id][-1]
                  < episode_1.partial_rewards_t[agent_id][-1]
              ):
                  indices_after_last_obs = episode_1.partial_rewards_t[
                      agent_id
                  ].find_indices_right(
                      episode_1.global_t_to_local_t[agent_id][-1]
                  )
                  episode_1_partial_rewards = list(
                      map(
                          episode_1.partial_rewards[
                              agent_id
                          ].__getitem__,
                          indices_after_last_obs,
                      )
                  )
                  self.assertEqual(
                      sum(episode_2.partial_rewards[agent_id]),
                      sum(episode_1_partial_rewards),
                  )
                  # Also ensure that the timestep mappings are correct.
                  episode_1_partial_rewards_t = list(
                      map(
                          episode_1.partial_rewards_t[
                              agent_id
                          ].__getitem__,
                          indices_after_last_obs,
                      )
                  )
                  self.assertListEqual(
                      episode_2.partial_rewards_t[agent_id],
                      episode_1_partial_rewards_t,
                  )
              # In the other case this mapping should be empty.
              else:
                  self.assertEqual(len(agent_global_ts), 0)
          # In the other case this mapping should be empty.
          else:
              self.assertEqual(len(agent_global_ts), 0)
  
      # Now test, if the specific values in the buffers are correct.
      (
          observations,
          actions,
          rewards,
          terminateds,
          truncateds,
          infos,
      ) = self._mock_multi_agent_records()
  
      # Create the episode.
      episode_1 = MultiAgentEpisode(
          agent_ids=[
              "agent_1",
              "agent_2",
              "agent_3",
              "agent_4",
              "agent_5",
          ],
          observations=observations,
          actions=actions,
          rewards=rewards,
          infos=infos,
          terminateds=terminateds,
          truncateds=truncateds,
      )
  
      # Assert that agents 1 and 3's buffers are indeed full.
      for agent_id in ["agent_1", "agent_3"]:
          self.assertEqual(
              actions[1][agent_id],
              episode_1.agent_buffers[agent_id]["actions"].queue[0],
          )
          # # Put the action back into the buffer.
          # episode_1.agent_buffers[agent_id]["actions"].put_nowait(
          #     actions[1][agent_id]
          # )
  
      # Now step once.
      action = {"agent_2": 3, "agent_4": 3}
      # This time agent 4 should have the buffer full, while agent 1 has emptied
      # its buffer.
      observation = {"agent_1": 3, "agent_2": 3}
      # Agents 1 and 2 add the reward to its timestep, but agent 3 and agent 5
      # add this to the buffer and to the global reward history.
      reward = {
          "agent_1": 2.0,
          "agent_2": 2.0,
          "agent_3": 2.0,
          "agent_5": 2.0,
      }
      info = {"agent_1": {}, "agent_2": {}}
      terminateds = {k: False for k in observation.keys()}
      terminateds.update({"__all__": False})
      truncateds = {k: False for k in observation.keys()}
      truncateds.update({"__all__": False})
      episode_1.add_env_step(
          observations=observation,
          actions=action,
          rewards=reward,
          infos=info,
          terminateds=terminateds,
          truncateds=truncateds,
      )
  
      # Check that the partial reward history is correct.
      self.assertEqual(len(episode_1.partial_rewards_t["agent_5"]), 1)
      self.assertEqual(len(episode_1.partial_rewards["agent_5"]), 1)
      self.assertEqual(len(episode_1.partial_rewards_t["agent_3"]), 2)
      self.assertEqual(len(episode_1.partial_rewards["agent_3"]), 2)
      self.assertEqual(len(episode_1.partial_rewards_t["agent_2"]), 2)
      self.assertEqual(len(episode_1.partial_rewards_t["agent_2"]), 2)
      self.assertListEqual(
          episode_1.partial_rewards["agent_3"], [0.5, 2.0]
      )
      self.assertListEqual(
          episode_1.partial_rewards_t["agent_3"], [1, 3]
      )
      self.assertListEqual(
          episode_1.partial_rewards["agent_2"], [1.0, 2.0]
      )
      self.assertListEqual(
          episode_1.partial_rewards_t["agent_2"], [2, 3]
      )
      self.assertEqual(len(episode_1.partial_rewards["agent_4"]), 2)
      self.assertListEqual(
          episode_1.partial_rewards["agent_4"], [0.5, 1.0]
      )
      self.assertListEqual(
          episode_1.partial_rewards_t["agent_4"], [1, 2]
      )
  
      # Now check that the reward buffers are full.
      for agent_id in ["agent_3", "agent_5"]:
          self.assertEqual(
              episode_1.agent_buffers[agent_id]["rewards"].queue[0], 2.0
          )
          # Check that the reward history is correctly recorded.
          self.assertEqual(
              episode_1.partial_rewards_t[agent_id][-1], episode_1.t
          )
          self.assertEqual(episode_1.partial_rewards[agent_id][-1], 2.0)
  
      # Now create the successor.
      episode_2 = episode_1.cut()
  
      for agent_id, agent_eps in episode_2.agent_episodes.items():
          if len(agent_eps.observations) > 0:
              # The successor's first observations should be the predecessor's last.
              self.assertEqual(
                  agent_eps.observations[0],
                  episode_1.agent_episodes[agent_id].observations[-1],
              )
              # The successor's first entry in the timestep mapping should be the
              # predecessor's last.
              self.assertEqual(
                  episode_2.global_t_to_local_t[agent_id][
                      -1
                  ],  # + episode_2.t_started,
                  episode_1.global_t_to_local_t[agent_id][-1],
              )
      # Now test that the partial rewards fit.
      for agent_id in ["agent_3", "agent_5"]:
          self.assertEqual(
              len(episode_2.partial_rewards_t[agent_id]), 1
          )
          self.assertEqual(episode_2.partial_rewards_t[agent_id][-1], 3)
          self.assertEqual(
              episode_2.agent_buffers[agent_id]["rewards"].queue[0], 2.0
          )
  
      # Assert that agent 3's and 4's action buffers are full.
      self.assertTrue(
          episode_2.agent_buffers["agent_4"]["actions"].full()
      )
      self.assertTrue(
          episode_2.agent_buffers["agent_3"]["actions"].full()
      )
      # Also assert that agent 1's action b uffer was emptied with the last
      # observations.
      self.assertTrue(
          episode_2.agent_buffers["agent_1"]["actions"].empty()
      )
  
imports: import gymnasium
  import numpy
  import unittest
  from typing import List, Optional, Tuple
  import ray
  from ray.rllib.env.multi_agent_env import MultiAgentEnv
  from ray.rllib.env.multi_agent_episode import MultiAgentEpisode
  from ray.rllib.utils.typing import MultiAgentDict
label_desc: Use image node_id 6 for calling the TestMultiAgentEpisode obj's underlying member method code with example usage: obj.test_cut() without return types