code: def test_getters(self):
      # Generate simple records for a multi agent environment.
      (
          observations,
          actions,
          rewards,
          is_terminateds,
          is_truncateds,
          infos,
      ) = self._mock_multi_agent_records()
      # Define the agent ids.
      agent_ids = [
          "agent_1",
          "agent_2",
          "agent_3",
          "agent_4",
          "agent_5",
      ]
  
      # Define some extra model outputs.
      extra_model_outputs = [
          # Here agent_2 has to buffer.
          {
              "agent_1": {"extra": 0},
              "agent_2": {"extra": 0},
              "agent_3": {"extra": 0},
          },
          {
              "agent_1": {"extra": 1},
              "agent_3": {"extra": 1},
              "agent_4": {"extra": 1},
          },
      ]
  
      # Create a multi-agent episode.
      episode = MultiAgentEpisode(
          agent_ids=agent_ids,
          observations=observations,
          actions=actions,
          rewards=rewards,
          infos=infos,
          terminateds=is_terminateds,
          truncateds=is_truncateds,
          extra_model_outputs=extra_model_outputs,
      )
  
      # --- observations ---
  
      # Now get the last observations.
      last_observation = episode.get_observations()
      # Assert that this observation is correct.
      for agent_id, agent_obs in last_observation.items():
          self.assertEqual(agent_obs[0], observations[-1][agent_id])
  
      last_observations = episode.get_observations(indices=[-1, -2])
      # Assert that the observations are correct.
      self.assertEqual(
          last_observations["agent_1"][0], observations[-2]["agent_1"]
      )
      self.assertEqual(
          last_observations["agent_2"][0], observations[-1]["agent_2"]
      )
      self.assertEqual(
          last_observations["agent_3"][0], observations[-2]["agent_3"]
      )
      # Note, agent 4 has two observations in the last two ones.
      # Note, `get_observations()` returns in the order of the `indices` parameter.
      self.assertEqual(
          last_observations["agent_4"][1], observations[-2]["agent_4"]
      )
      self.assertEqual(
          last_observations["agent_4"][0], observations[-1]["agent_4"]
      )
  
      # Now, test the same when returning a list.
      last_observation = episode.get_observations(as_list=True)
      # Assert that these observations are correct.
      for agent_id, agent_obs in last_observation[0].items():
          self.assertEqual(agent_obs, observations[-1][agent_id])
  
      last_observations = episode.get_observations(
          indices=[-1, -2], as_list=True
      )
      # Assert that the observations are correct.
      self.assertEqual(
          last_observations[1]["agent_1"], observations[-2]["agent_1"]
      )
      self.assertEqual(
          last_observations[0]["agent_2"], observations[-1]["agent_2"]
      )
      self.assertEqual(
          last_observations[1]["agent_3"], observations[-2]["agent_3"]
      )
      # Note, agent 4 has two observations in the last two ones.
      # Note, `get_observations()` returns in the order of the `indices` parameter.
      self.assertEqual(
          last_observations[1]["agent_4"], observations[-2]["agent_4"]
      )
      self.assertEqual(
          last_observations[0]["agent_4"], observations[-1]["agent_4"]
      )
  
      # Now, test if for the local setting the results fit as well.
      last_local_observation = episode.get_observations(global_ts=False)
      # Assert that the observations are correct.
      self.assertEqual(
          last_local_observation["agent_1"][0],
          observations[-2]["agent_1"],
      )
      self.assertEqual(
          last_local_observation["agent_2"][0],
          observations[-1]["agent_2"],
      )
      self.assertEqual(
          last_local_observation["agent_3"][0],
          observations[-2]["agent_3"],
      )
      self.assertEqual(
          last_local_observation["agent_4"][0],
          observations[-1]["agent_4"],
      )
  
      # Now return the last two observations per agent.
      last_local_observations = episode.get_observations(
          indices=[-1, -2], global_ts=False
      )
      # Assert that the observations are correct.
      self.assertEqual(
          last_local_observations["agent_1"][0],
          observations[-2]["agent_1"],
      )
      self.assertEqual(
          last_local_observations["agent_1"][1],
          observations[-3]["agent_1"],
      )
      self.assertEqual(
          last_local_observations["agent_2"][0],
          observations[-1]["agent_2"],
      )
      self.assertEqual(
          last_local_observations["agent_2"][1],
          observations[-3]["agent_2"],
      )
      self.assertEqual(
          last_local_observations["agent_3"][0],
          observations[-2]["agent_3"],
      )
      self.assertEqual(
          last_local_observations["agent_3"][1],
          observations[-3]["agent_3"],
      )
      self.assertEqual(
          last_local_observations["agent_4"][0],
          observations[-1]["agent_4"],
      )
      self.assertEqual(
          last_local_observations["agent_4"][1],
          observations[-2]["agent_4"],
      )
  
      # Test with initial observations only.
      episode_init_only = MultiAgentEpisode(agent_ids=agent_ids)
      episode_init_only.add_env_reset(
          observations=observations[0],
          infos=infos[0],
      )
      # Get the last observation for agents and assert that its correct.
      last_observation = episode_init_only.get_observations()
      for agent_id, agent_obs in observations[0].items():
          self.assertEqual(last_observation[agent_id][0], agent_obs)
      # Now the same as list.
      last_observation = episode_init_only.get_observations(
          as_list=True
      )
      for agent_id, agent_obs in observations[0].items():
          self.assertEqual(last_observation[0][agent_id], agent_obs)
      # Now locally.
      last_local_observation = episode_init_only.get_observations(
          global_ts=False
      )
      for agent_id, agent_obs in observations[0].items():
          self.assertEqual(
              last_local_observation[agent_id][0], agent_obs
          )
  
      # --- actions ---
      last_actions = episode.get_actions()
      self.assertEqual(
          last_actions["agent_1"][0], actions[-1]["agent_1"]
      )
      self.assertEqual(
          last_actions["agent_3"][0], actions[-1]["agent_3"]
      )
      self.assertEqual(
          last_actions["agent_4"][0], actions[-1]["agent_4"]
      )
  
      last_actions = episode.get_actions(indices=[-1, -2])
      self.assertEqual(
          last_actions["agent_1"][0], actions[-1]["agent_1"]
      )
      self.assertEqual(
          last_actions["agent_3"][0], actions[-1]["agent_3"]
      )
      self.assertEqual(
          last_actions["agent_4"][0], actions[-1]["agent_4"]
      )
      self.assertEqual(
          last_actions["agent_1"][1], actions[-2]["agent_1"]
      )
      self.assertEqual(
          last_actions["agent_2"][0], actions[-2]["agent_2"]
      )
      self.assertEqual(
          last_actions["agent_3"][1], actions[-2]["agent_3"]
      )
  
      # Now request lists.
      last_actions = episode.get_actions(as_list=True)
      self.assertEqual(
          last_actions[0]["agent_1"], actions[-1]["agent_1"]
      )
      self.assertEqual(
          last_actions[0]["agent_3"], actions[-1]["agent_3"]
      )
      self.assertEqual(
          last_actions[0]["agent_4"], actions[-1]["agent_4"]
      )
  
      # Request the last two actions and return as a list.
      last_actions = episode.get_actions([-1, -2], as_list=True)
      self.assertEqual(
          last_actions[0]["agent_1"], actions[-1]["agent_1"]
      )
      self.assertEqual(
          last_actions[0]["agent_3"], actions[-1]["agent_3"]
      )
      self.assertEqual(
          last_actions[0]["agent_4"], actions[-1]["agent_4"]
      )
      self.assertEqual(
          last_actions[1]["agent_1"], actions[-2]["agent_1"]
      )
      self.assertEqual(
          last_actions[1]["agent_2"], actions[-2]["agent_2"]
      )
      self.assertEqual(
          last_actions[1]["agent_3"], actions[-2]["agent_3"]
      )
  
      # Now request the last actions at the local timesteps, i.e. for each agent
      # its last two actions.
      last_actions = episode.get_actions([-1, -2], global_ts=False)
      self.assertEqual(
          last_actions["agent_1"][0], actions[-1]["agent_1"]
      )
      self.assertEqual(
          last_actions["agent_3"][0], actions[-1]["agent_3"]
      )
      self.assertEqual(
          last_actions["agent_4"][0], actions[-1]["agent_4"]
      )
      self.assertEqual(
          last_actions["agent_1"][1], actions[-2]["agent_1"]
      )
      self.assertEqual(
          last_actions["agent_2"][0], actions[-2]["agent_2"]
      )
      self.assertEqual(
          last_actions["agent_3"][1], actions[-2]["agent_3"]
      )
  
      # --- extra_model_outputs ---
      last_extra_model_outputs = episode.get_extra_model_outputs(
          "extra"
      )
      self.assertEqual(
          last_extra_model_outputs["agent_1"][0],
          extra_model_outputs[-1]["agent_1"]["extra"],
      )
      self.assertEqual(
          last_extra_model_outputs["agent_3"][0],
          extra_model_outputs[-1]["agent_3"]["extra"],
      )
      self.assertEqual(
          last_extra_model_outputs["agent_4"][0],
          extra_model_outputs[-1]["agent_4"]["extra"],
      )
  
      # Request the last two outputs.
      last_extra_model_outputs = episode.get_extra_model_outputs(
          "extra", indices=[-1, -2]
      )
      self.assertEqual(
          last_extra_model_outputs["agent_1"][0],
          extra_model_outputs[-1]["agent_1"]["extra"],
      )
      self.assertEqual(
          last_extra_model_outputs["agent_3"][0],
          extra_model_outputs[-1]["agent_3"]["extra"],
      )
      self.assertEqual(
          last_extra_model_outputs["agent_4"][0],
          extra_model_outputs[-1]["agent_4"]["extra"],
      )
      self.assertEqual(
          last_extra_model_outputs["agent_1"][1],
          extra_model_outputs[-2]["agent_1"]["extra"],
      )
      self.assertEqual(
          last_extra_model_outputs["agent_2"][0],
          extra_model_outputs[-2]["agent_2"]["extra"],
      )
      self.assertEqual(
          last_extra_model_outputs["agent_3"][1],
          extra_model_outputs[-2]["agent_3"]["extra"],
      )
  
      # Now request lists.
      last_extra_model_outputs = episode.get_extra_model_outputs(
          "extra", as_list=True
      )
      self.assertEqual(
          last_extra_model_outputs[0]["agent_1"],
          extra_model_outputs[-1]["agent_1"]["extra"],
      )
      self.assertEqual(
          last_extra_model_outputs[0]["agent_3"],
          extra_model_outputs[-1]["agent_3"]["extra"],
      )
      self.assertEqual(
          last_extra_model_outputs[0]["agent_4"],
          extra_model_outputs[-1]["agent_4"]["extra"],
      )
      # Request the last two extra model outputs and return as a list.
      last_extra_model_outputs = episode.get_extra_model_outputs(
          "extra", [-1, -2], as_list=True
      )
      self.assertEqual(
          last_extra_model_outputs[0]["agent_1"],
          extra_model_outputs[-1]["agent_1"]["extra"],
      )
      self.assertEqual(
          last_extra_model_outputs[0]["agent_3"],
          extra_model_outputs[-1]["agent_3"]["extra"],
      )
      self.assertEqual(
          last_extra_model_outputs[0]["agent_4"],
          extra_model_outputs[-1]["agent_4"]["extra"],
      )
      self.assertEqual(
          last_extra_model_outputs[1]["agent_1"],
          extra_model_outputs[-2]["agent_1"]["extra"],
      )
      self.assertEqual(
          last_extra_model_outputs[1]["agent_2"],
          extra_model_outputs[-2]["agent_2"]["extra"],
      )
      self.assertEqual(
          last_extra_model_outputs[1]["agent_3"],
          extra_model_outputs[-2]["agent_3"]["extra"],
      )
  
      # Now request the last extra model outputs at the local timesteps, i.e.
      # for each agent its last two actions.
      last_extra_model_outputs = episode.get_extra_model_outputs(
          "extra", [-1, -2], global_ts=False
      )
      self.assertEqual(
          last_extra_model_outputs["agent_1"][0],
          extra_model_outputs[-1]["agent_1"]["extra"],
      )
      self.assertEqual(
          last_extra_model_outputs["agent_3"][0],
          extra_model_outputs[-1]["agent_3"]["extra"],
      )
      self.assertEqual(
          last_extra_model_outputs["agent_4"][0],
          extra_model_outputs[-1]["agent_4"]["extra"],
      )
      self.assertEqual(
          last_extra_model_outputs["agent_1"][1],
          extra_model_outputs[-2]["agent_1"]["extra"],
      )
      self.assertEqual(
          last_extra_model_outputs["agent_2"][0],
          extra_model_outputs[-2]["agent_2"]["extra"],
      )
      self.assertEqual(
          last_extra_model_outputs["agent_3"][1],
          extra_model_outputs[-2]["agent_3"]["extra"],
      )
  
      # TODO (simon): Not tested with `global_ts=False`.
      # --- rewards ---
      # Start with the case of no partial or buffered rewards.
      last_rewards = episode.get_rewards(
          partial=False, consider_buffer=False
      )
      self.assertTrue(
          last_rewards["agent_4"][0],
          rewards[0]["agent_4"] + rewards[1]["agent_4"],
      )
      self.assertTrue(last_rewards["agent_2"][0], rewards[1]["agent_2"])
  
      # Now test the same case, but with the last two rewards.
      last_rewards = episode.get_rewards(
          [-1, -2], partial=False, consider_buffer=False
      )
      self.assertTrue(
          last_rewards["agent_4"][0],
          rewards[0]["agent_4"] + rewards[1]["agent_4"],
      )
      self.assertTrue(last_rewards["agent_2"][0], rewards[1]["agent_2"])
      self.assertTrue(last_rewards["agent_1"][0], rewards[0]["agent_1"])
      self.assertTrue(last_rewards["agent_3"][0], rewards[0]["agent_3"])
  
      # Now request these rewards as list.
      last_rewards = episode.get_rewards(
          as_list=True, partial=False, consider_buffer=False
      )
      self.assertTrue(
          last_rewards[0]["agent_4"],
          rewards[0]["agent_4"] + rewards[1]["agent_4"],
      )
      self.assertTrue(last_rewards[0]["agent_2"], rewards[1]["agent_2"])
  
      # Now test the same case, but with the last two rewards.
      last_rewards = episode.get_rewards(
          [-1, -2], as_list=True, partial=False, consider_buffer=False
      )
      self.assertTrue(
          last_rewards[0]["agent_4"],
          rewards[0]["agent_4"] + rewards[1]["agent_4"],
      )
      self.assertTrue(last_rewards[0]["agent_2"], rewards[1]["agent_2"])
      self.assertTrue(last_rewards[1]["agent_1"], rewards[0]["agent_1"])
      self.assertTrue(last_rewards[1]["agent_3"], rewards[0]["agent_3"])
  
      # Create an environment.
      env = MultiAgentTestEnv()
      # Create an empty episode.
      episode_1 = MultiAgentEpisode(agent_ids=env._agent_ids)
  
      # Generate initial observation and info.
      obs, info = env.reset(seed=42)
      episode_1.add_env_reset(
          observations=obs,
          infos=info,
      )
      # Now, generate 100 samples.
      for i in range(100):
          action = {agent_id: i for agent_id in obs}
          obs, reward, terminated, truncated, info = env.step(action)
          episode_1.add_env_step(
              observations=obs,
              actions=action,
              rewards=reward,
              infos=info,
              terminateds=terminated,
              truncateds=truncated,
              extra_model_outputs={
                  agent_id: {"extra": 10} for agent_id in action
              },
          )
  
      # First, receive the last rewards without considering buffered values.
      last_rewards = episode_1.get_rewards(
          partial=False, consider_buffer=False
      )
      self.assertIn("agent_9", last_rewards)
      self.assertEqual(
          episode_1.global_t_to_local_t["agent_9"][-1], 100
      )
      self.assertEqual(
          episode_1.agent_episodes["agent_9"].rewards[-1], 1.0
      )
      self.assertEqual(last_rewards["agent_9"][0], 1.0)
      self.assertIn("agent_0", last_rewards)
      self.assertEqual(
          episode_1.global_t_to_local_t["agent_0"][-1], 100
      )
      self.assertEqual(
          episode_1.agent_episodes["agent_0"].rewards[-1], 1.0
      )
      self.assertEqual(last_rewards["agent_0"][0], 1.0)
      self.assertIn("agent_2", last_rewards)
      self.assertEqual(
          episode_1.global_t_to_local_t["agent_2"][-1], 100
      )
      self.assertEqual(
          episode_1.agent_episodes["agent_2"].rewards[-1], 1.0
      )
      self.assertEqual(last_rewards["agent_2"][0], 1.0)
      self.assertIn("agent_5", last_rewards)
      self.assertEqual(
          episode_1.global_t_to_local_t["agent_5"][-1], 100
      )
      self.assertEqual(
          episode_1.agent_episodes["agent_5"].rewards[-1], 1.0
      )
      self.assertEqual(last_rewards["agent_5"][0], 1.0)
      self.assertIn("agent_8", last_rewards)
      self.assertEqual(
          episode_1.global_t_to_local_t["agent_8"][-1], 100
      )
      self.assertEqual(
          episode_1.agent_episodes["agent_8"].rewards[-1], 1.0
      )
      self.assertEqual(last_rewards["agent_8"][0], 1.0)
      self.assertIn("agent_4", last_rewards)
      self.assertEqual(
          episode_1.global_t_to_local_t["agent_4"][-1], 100
      )
      self.assertEqual(
          episode_1.agent_episodes["agent_4"].rewards[-1], 1.0
      )
      self.assertEqual(last_rewards["agent_4"][0], 1.0)
      self.assertIn("agent_3", last_rewards)
      self.assertEqual(
          episode_1.global_t_to_local_t["agent_3"][-1], 100
      )
      # Agent 3 had a partial reward before the last recorded observation.
      self.assertEqual(
          episode_1.agent_episodes["agent_3"].rewards[-1], 2.0
      )
      self.assertEqual(last_rewards["agent_3"][0], 2.0)
      # Assert that all the other agents are not in the returned rewards.
      self.assertNotIn("agent_1", last_rewards)
      self.assertNotIn("agent_6", last_rewards)
      self.assertNotIn("agent_7", last_rewards)
  
      # Now return the same as list.
      last_rewards = episode_1.get_rewards(
          partial=False, consider_buffer=False, as_list=True
      )
      self.assertIn("agent_9", last_rewards[0])
      self.assertEqual(last_rewards[0]["agent_9"], 1.0)
      self.assertIn("agent_0", last_rewards[0])
      self.assertEqual(last_rewards[0]["agent_0"], 1.0)
      self.assertIn("agent_2", last_rewards[0])
      self.assertEqual(last_rewards[0]["agent_2"], 1.0)
      self.assertIn("agent_5", last_rewards[0])
      self.assertEqual(last_rewards[0]["agent_5"], 1.0)
      self.assertIn("agent_8", last_rewards[0])
      self.assertEqual(last_rewards[0]["agent_8"], 1.0)
      self.assertIn("agent_4", last_rewards[0])
      self.assertEqual(last_rewards[0]["agent_4"], 1.0)
      self.assertIn("agent_3", last_rewards[0])
      self.assertEqual(last_rewards[0]["agent_3"], 2.0)
      # Assert that all the other agents are not in the returned rewards.
      self.assertNotIn("agent_1", last_rewards)
      self.assertNotIn("agent_6", last_rewards)
      self.assertNotIn("agent_7", last_rewards)
  
      # Now request the last two indices.
      last_rewards = episode_1.get_rewards(
          [-1, -2], partial=False, consider_buffer=False
      )
      self.assertIn("agent_9", last_rewards)
      self.assertListEqual(
          episode_1.global_t_to_local_t["agent_9"][-2:], [99, 100]
      )
      self.assertListEqual(
          episode_1.agent_episodes["agent_9"].rewards[-2:],
          last_rewards["agent_9"],
      )
      self.assertIn("agent_5", last_rewards)
      self.assertListEqual(
          episode_1.global_t_to_local_t["agent_5"][-2:], [99, 100]
      )
      # Agent 5 has already died, so we need to convert back to list.
      self.assertListEqual(
          episode_1.agent_episodes["agent_5"].rewards[-2:],
          last_rewards["agent_5"],
      )
      self.assertIn("agent_2", last_rewards)
      self.assertListEqual(
          episode_1.global_t_to_local_t["agent_2"][-2:], [99, 100]
      )
      self.assertListEqual(
          episode_1.agent_episodes["agent_2"].rewards[-1:-3:-1],
          last_rewards["agent_2"],
      )
      # Agent 2 had no observation at `ts=98`, but partial rewards.
      self.assertGreater(
          99, episode_1.global_t_to_local_t["agent_2"][-3]
      )
      # Ensure that for agent 2 there had been three partial rewards in between the
      # observation at `ts=95` and the next at `ts=99`.
      self.assertListEqual(
          episode_1.partial_rewards_t["agent_2"][-4:-1], [96, 98, 99]
      )
      self.assertIn("agent_3", last_rewards)
      # Agent 3 had no observation at `ts=99`.
      self.assertListEqual(
          episode_1.global_t_to_local_t["agent_3"][-2:], [98, 100]
      )
      self.assertEqual(
          episode_1.agent_episodes["agent_3"].rewards[-1],
          last_rewards["agent_3"][0],
      )
      # Ensure that there was a partial reward at `ts=99`.
      self.assertListEqual(
          episode_1.partial_rewards_t["agent_3"][-2:], [99, 100]
      )
      self.assertIn("agent_4", last_rewards)
      self.assertListEqual(
          episode_1.global_t_to_local_t["agent_4"][-2:], [99, 100]
      )
      self.assertListEqual(
          episode_1.agent_episodes["agent_4"].rewards[-2:],
          last_rewards["agent_4"],
      )
      self.assertIn("agent_8", last_rewards)
      # Ensure that the third-last observation is before `ts=98`.
      self.assertListEqual(
          episode_1.global_t_to_local_t["agent_8"][-3:], [97, 99, 100]
      )
      # Ensure also that at `ts=97` there was a reward.
      self.assertListEqual(
          episode_1.partial_rewards_t["agent_8"][-3:-1], [98, 99]
      )
      self.assertListEqual([1.0, 2.0], last_rewards["agent_8"])
      self.assertIn("agent_7", last_rewards)
      # Agent 7 has no observation at `ts=100`, but at `ts=98`.
      self.assertListEqual(
          episode_1.global_t_to_local_t["agent_7"][-2:], [98, 99]
      )
      self.assertEqual(
          episode_1.agent_episodes["agent_7"].rewards[-1],
          last_rewards["agent_7"][0],
      )
      self.assertIn("agent_0", last_rewards)
      self.assertListEqual(
          episode_1.global_t_to_local_t["agent_0"][-2:], [99, 100]
      )
      self.assertListEqual(
          episode_1.agent_episodes["agent_0"].rewards[-2:],
          last_rewards["agent_0"],
      )
      self.assertNotIn("agent_1", last_rewards)
      self.assertNotIn("agent_6", last_rewards)
  
      # Now request the last two indices as list.
      last_rewards = episode_1.get_rewards(
          [-1, -2], partial=False, consider_buffer=False, as_list=True
      )
      self.assertIn("agent_9", last_rewards[0])
      self.assertIn("agent_9", last_rewards[1])
      self.assertEqual(
          episode_1.agent_episodes["agent_9"].rewards[-1],
          last_rewards[0]["agent_9"],
      )
      self.assertEqual(
          episode_1.agent_episodes["agent_9"].rewards[-2],
          last_rewards[1]["agent_9"],
      )
      self.assertIn("agent_5", last_rewards[0])
      self.assertIn("agent_5", last_rewards[1])
      self.assertEqual(
          episode_1.agent_episodes["agent_5"].rewards[-1],
          last_rewards[0]["agent_5"],
      )
      self.assertEqual(
          episode_1.agent_episodes["agent_5"].rewards[-2],
          last_rewards[1]["agent_5"],
      )
      self.assertIn("agent_2", last_rewards[0])
      self.assertIn("agent_2", last_rewards[1])
      self.assertEqual(
          episode_1.agent_episodes["agent_2"].rewards[-1],
          last_rewards[0]["agent_2"],
      )
      self.assertEqual(3.0, last_rewards[1]["agent_2"])
      # Agent 3 has only recorded rewards at `ts=100`.
      self.assertIn("agent_3", last_rewards[0])
      self.assertEqual(
          episode_1.agent_episodes["agent_3"].rewards[-1],
          last_rewards[0]["agent_3"],
      )
      self.assertIn("agent_4", last_rewards[0])
      self.assertIn("agent_4", last_rewards[1])
      self.assertEqual(
          episode_1.agent_episodes["agent_4"].rewards[-1],
          last_rewards[0]["agent_4"],
      )
      self.assertEqual(
          episode_1.agent_episodes["agent_4"].rewards[-2],
          last_rewards[1]["agent_4"],
      )
      self.assertIn("agent_8", last_rewards[0])
      self.assertIn("agent_8", last_rewards[1])
      self.assertEqual(
          episode_1.agent_episodes["agent_8"].rewards[-1],
          last_rewards[0]["agent_8"],
      )
      self.assertEqual(
          episode_1.agent_episodes["agent_8"].rewards[-2],
          last_rewards[1]["agent_8"],
      )
      # Agent 7 has no observation at `ts=100`.
      self.assertIn("agent_7", last_rewards[1])
      self.assertEqual(
          episode_1.agent_episodes["agent_7"].rewards[-1],
          last_rewards[1]["agent_7"],
      )
      self.assertIn("agent_0", last_rewards[0])
      self.assertIn("agent_0", last_rewards[1])
      self.assertEqual(
          episode_1.agent_episodes["agent_0"].rewards[-1],
          last_rewards[0]["agent_0"],
      )
      self.assertEqual(
          episode_1.agent_episodes["agent_0"].rewards[-2],
          last_rewards[1]["agent_0"],
      )
      self.assertNotIn("agent_1", last_rewards[0])
      self.assertNotIn("agent_6", last_rewards[0])
      self.assertNotIn("agent_1", last_rewards[1])
      self.assertNotIn("agent_6", last_rewards[1])
  
      # Second, get the last rewards with a single index, consider all partial
      # rewards after the last recorded observation of an agent, i.e. set
      # `consider_buffer` to `True`.
      last_rewards = episode_1.get_rewards(
          partial=False, consider_buffer=True
      )
      self.assertIn("agent_9", last_rewards)
      self.assertEqual(
          episode_1.agent_episodes["agent_9"].rewards[-1],
          last_rewards["agent_9"][0],
      )
      self.assertIn("agent_0", last_rewards)
      self.assertEqual(
          episode_1.agent_episodes["agent_0"].rewards[-1],
          last_rewards["agent_0"][0],
      )
      self.assertIn("agent_2", last_rewards)
      self.assertEqual(
          episode_1.agent_episodes["agent_2"].rewards[-1],
          last_rewards["agent_2"][0],
      )
      self.assertIn("agent_5", last_rewards)
      self.assertEqual(
          episode_1.agent_episodes["agent_5"].rewards[-1],
          last_rewards["agent_5"][0],
      )
      self.assertIn("agent_8", last_rewards)
      self.assertEqual(
          episode_1.agent_episodes["agent_8"].rewards[-1],
          last_rewards["agent_8"][0],
      )
      self.assertIn("agent_4", last_rewards)
      self.assertEqual(
          episode_1.agent_episodes["agent_4"].rewards[-1],
          last_rewards["agent_4"][0],
      )
      self.assertIn("agent_3", last_rewards)
      # Agent 3 had a partial reward before the last recorded observation.
      self.assertEqual(
          episode_1.agent_episodes["agent_3"].rewards[-1],
          last_rewards["agent_3"][0],
      )
      # Agent 7 has a partial reward at `ts=100` after its last observation at
      # `ts=99`.
      self.assertIn("agent_7", last_rewards)
      self.assertEqual(episode_1.partial_rewards_t["agent_7"][-1], 100)
      self.assertEqual(
          episode_1.partial_rewards["agent_7"][-1],
          last_rewards["agent_7"][0],
      )
      # Assert that all the other agents are not in the returned rewards.
      self.assertNotIn("agent_1", last_rewards)
      self.assertNotIn("agent_6", last_rewards)
  
      # Now request the last rewards as a list while considering the buffer.
      last_rewards = episode_1.get_rewards(
          partial=False, consider_buffer=True, as_list=True
      )
      self.assertIn("agent_9", last_rewards[0])
      self.assertEqual(
          episode_1.agent_episodes["agent_9"].rewards[-1],
          last_rewards[0]["agent_9"],
      )
      self.assertIn("agent_0", last_rewards[0])
      self.assertEqual(
          episode_1.agent_episodes["agent_0"].rewards[-1],
          last_rewards[0]["agent_0"],
      )
      self.assertIn("agent_2", last_rewards[0])
      self.assertEqual(
          episode_1.agent_episodes["agent_2"].rewards[-1],
          last_rewards[0]["agent_2"],
      )
      self.assertIn("agent_5", last_rewards[0])
      self.assertEqual(
          episode_1.agent_episodes["agent_5"].rewards[-1],
          last_rewards[0]["agent_5"],
      )
      self.assertIn("agent_8", last_rewards[0])
      self.assertEqual(
          episode_1.agent_episodes["agent_8"].rewards[-1],
          last_rewards[0]["agent_8"],
      )
      self.assertIn("agent_4", last_rewards[0])
      self.assertEqual(
          episode_1.agent_episodes["agent_4"].rewards[-1],
          last_rewards[0]["agent_4"],
      )
      self.assertIn("agent_3", last_rewards[0])
      # Agent 3 had a partial reward before the last recorded observation.
      self.assertEqual(
          episode_1.agent_episodes["agent_3"].rewards[-1],
          last_rewards[0]["agent_3"],
      )
      # Agent 7 has a partial reward at `ts=100` after its last observation at
      # `ts=99`.
      self.assertIn("agent_7", last_rewards[0])
      self.assertEqual(
          episode_1.partial_rewards["agent_7"][-1],
          last_rewards[0]["agent_7"],
      )
      # Assert that all the other agents are not in the returned rewards.
      self.assertNotIn("agent_1", last_rewards[0])
      self.assertNotIn("agent_6", last_rewards[0])
  
      # Now request the last two indices and consider buffered partial rewards after
      # the last observation.
      last_rewards = episode_1.get_rewards(
          [-1, -2], partial=False, consider_buffer=True
      )
      self.assertIn("agent_9", last_rewards)
      self.assertListEqual(
          episode_1.agent_episodes["agent_9"].rewards[-1:-3:-1],
          last_rewards["agent_9"],
      )
      self.assertIn("agent_0", last_rewards)
      self.assertListEqual(
          episode_1.agent_episodes["agent_0"].rewards[-1:-3:-1],
          last_rewards["agent_0"],
      )
      self.assertIn("agent_2", last_rewards)
      self.assertListEqual(
          episode_1.agent_episodes["agent_2"].rewards[-1:-3:-1],
          last_rewards["agent_2"],
      )
      self.assertIn("agent_5", last_rewards)
      # Agent 5 already died, so we need to convert to list first.
      self.assertListEqual(
          episode_1.agent_episodes["agent_5"].rewards[-1:-3:-1],
          last_rewards["agent_5"],
      )
      self.assertIn("agent_8", last_rewards)
      self.assertListEqual(
          episode_1.agent_episodes["agent_8"].rewards[-1:-3:-1],
          last_rewards["agent_8"],
      )
      self.assertIn("agent_4", last_rewards)
      self.assertListEqual(
          episode_1.agent_episodes["agent_4"].rewards[-1:-3:-1],
          last_rewards["agent_4"],
      )
      # Nothing changes for agent 3 as it has an observation at the last requested
      # timestep 100, but not at `ts=99`.
      self.assertIn("agent_3", last_rewards)
      self.assertEqual(
          episode_1.agent_episodes["agent_3"].rewards[-1],
          last_rewards["agent_3"][0],
      )
      # The entries for agent 6 have changed now b/c it has partial rewards during the
      # requested timesteps 100 and 99.
      self.assertIn("agent_6", last_rewards)
      self.assertListEqual(
          episode_1.global_t_to_local_t["agent_6"][-2:], [95, 98]
      )
      self.assertListEqual(
          episode_1.partial_rewards_t["agent_6"][-2:], [99, 100]
      )
      self.assertListEqual(
          episode_1.partial_rewards["agent_6"][-2:],
          last_rewards["agent_6"],
      )
      # Entries for agent 7 also change b/c this agent has a partial reward at
      # `ts=100` while it has no observation recorded at this timestep.
      self.assertIn("agent_7", last_rewards)
      self.assertListEqual(
          episode_1.global_t_to_local_t["agent_7"][-2:], [98, 99]
      )
      self.assertListEqual(
          episode_1.partial_rewards_t["agent_7"][-2:], [99, 100]
      )
      self.assertEqual(
          episode_1.partial_rewards["agent_7"][-1],
          last_rewards["agent_7"][0],
      )
      self.assertEqual(
          episode_1.agent_episodes["agent_7"].rewards[-1],
          last_rewards["agent_7"][1],
      )
      # Assert that all the other agents are not in the returned rewards.
      self.assertNotIn("agent_1", last_rewards)
  
      # Now request the same indices with `consider_buffer=True` and return them as
      # a list.
      last_rewards = episode_1.get_rewards(
          [-1, -2], partial=False, consider_buffer=True, as_list=True
      )
      self.assertIn("agent_9", last_rewards[0])
      self.assertIn("agent_9", last_rewards[1])
      self.assertEqual(
          episode_1.agent_episodes["agent_9"].rewards[-1],
          last_rewards[0]["agent_9"],
      )
      self.assertEqual(
          episode_1.agent_episodes["agent_9"].rewards[-2],
          last_rewards[1]["agent_9"],
      )
      self.assertIn("agent_0", last_rewards[0])
      self.assertIn("agent_0", last_rewards[1])
      self.assertEqual(
          episode_1.agent_episodes["agent_0"].rewards[-1],
          last_rewards[0]["agent_0"],
      )
      self.assertEqual(
          episode_1.agent_episodes["agent_0"].rewards[-2],
          last_rewards[1]["agent_0"],
      )
      self.assertIn("agent_2", last_rewards[0])
      self.assertIn("agent_2", last_rewards[1])
      self.assertEqual(
          episode_1.agent_episodes["agent_2"].rewards[-1],
          last_rewards[0]["agent_2"],
      )
      self.assertEqual(
          episode_1.agent_episodes["agent_2"].rewards[-2],
          last_rewards[1]["agent_2"],
      )
      self.assertIn("agent_5", last_rewards[0])
      self.assertIn("agent_5", last_rewards[1])
      self.assertEqual(
          episode_1.agent_episodes["agent_5"].rewards[-1],
          last_rewards[0]["agent_5"],
      )
      self.assertEqual(
          episode_1.agent_episodes["agent_5"].rewards[-2],
          last_rewards[1]["agent_5"],
      )
      self.assertIn("agent_8", last_rewards[0])
      self.assertIn("agent_8", last_rewards[1])
      self.assertEqual(
          episode_1.agent_episodes["agent_8"].rewards[-1],
          last_rewards[0]["agent_8"],
      )
      self.assertEqual(
          episode_1.agent_episodes["agent_8"].rewards[-2],
          last_rewards[1]["agent_8"],
      )
      self.assertIn("agent_4", last_rewards[0])
      self.assertIn("agent_4", last_rewards[1])
      self.assertEqual(
          episode_1.agent_episodes["agent_4"].rewards[-1],
          last_rewards[0]["agent_4"],
      )
      self.assertEqual(
          episode_1.agent_episodes["agent_4"].rewards[-2],
          last_rewards[1]["agent_4"],
      )
      # Nothing changes for agent 3 as it has an observation at the last requested
      # timestep 100.
      self.assertIn("agent_3", last_rewards[0])
      self.assertNotIn("agent_3", last_rewards[1])
      self.assertEqual(
          episode_1.agent_episodes["agent_3"].rewards[-1],
          last_rewards[0]["agent_3"],
      )
      # The entries for agent 6 have changed now b/c it has partial rewards during the
      # requested timesteps 100 and 99.
      self.assertIn("agent_6", last_rewards[0])
      self.assertIn("agent_6", last_rewards[1])
      self.assertEqual(
          episode_1.partial_rewards["agent_6"][-1],
          last_rewards[0]["agent_6"],
      )
      self.assertEqual(
          episode_1.partial_rewards["agent_6"][-2],
          last_rewards[1]["agent_6"],
      )
      # Entries for agent 7 also change b/c this agent has a partial reward at
      # `ts=100` while it has no observation recorded at this timestep.
      self.assertIn("agent_7", last_rewards[0])
      self.assertIn("agent_7", last_rewards[1])
      self.assertEqual(
          episode_1.partial_rewards["agent_7"][-1],
          last_rewards[0]["agent_7"],
      )
      self.assertEqual(
          episode_1.agent_episodes["agent_7"].rewards[-1],
          last_rewards[1]["agent_7"],
      )
      # Assert that all the other agents are not in the returned rewards.
      self.assertNotIn("agent_1", last_rewards[0])
      self.assertNotIn("agent_1", last_rewards[1])
  
      # Third, request only partial rewards, i.e. rewards do not get buffered and
      # added up.
      last_rewards = episode_1.get_rewards(
          partial=True, consider_buffer=False
      )
      self.assertIn("agent_9", last_rewards)
      self.assertEqual(episode_1.partial_rewards_t["agent_9"][-1], 100)
      self.assertEqual(
          episode_1.partial_rewards["agent_9"][-1],
          last_rewards["agent_9"][-1],
      )
      self.assertIn("agent_0", last_rewards)
      self.assertEqual(episode_1.partial_rewards_t["agent_0"][-1], 100)
      self.assertEqual(
          episode_1.partial_rewards["agent_0"][-1],
          last_rewards["agent_0"][-1],
      )
      self.assertIn("agent_2", last_rewards)
      self.assertEqual(episode_1.partial_rewards_t["agent_2"][-1], 100)
      self.assertEqual(
          episode_1.partial_rewards["agent_2"][-1],
          last_rewards["agent_2"][-1],
      )
      self.assertIn("agent_8", last_rewards)
      self.assertEqual(episode_1.partial_rewards_t["agent_8"][-1], 100)
      self.assertEqual(
          episode_1.partial_rewards["agent_8"][-1],
          last_rewards["agent_8"][-1],
      )
      self.assertIn("agent_4", last_rewards)
      self.assertEqual(episode_1.partial_rewards_t["agent_4"][-1], 100)
      self.assertEqual(
          episode_1.partial_rewards["agent_4"][-1],
          last_rewards["agent_4"][-1],
      )
      self.assertIn("agent_3", last_rewards)
      self.assertEqual(episode_1.partial_rewards_t["agent_3"][-1], 100)
      self.assertEqual(
          episode_1.partial_rewards["agent_3"][-1],
          last_rewards["agent_3"][-1],
      )
      self.assertIn("agent_6", last_rewards)
      self.assertEqual(episode_1.partial_rewards_t["agent_6"][-1], 100)
      self.assertEqual(
          episode_1.partial_rewards["agent_6"][-1],
          last_rewards["agent_6"][-1],
      )
      self.assertIn("agent_7", last_rewards)
      self.assertEqual(episode_1.partial_rewards_t["agent_7"][-1], 100)
      self.assertEqual(
          episode_1.partial_rewards["agent_7"][-1],
          last_rewards["agent_7"][-1],
      )
      # Assert that all the other agents are not in the returned rewards.
      self.assertNotIn("agent_1", last_rewards)
  
      # Now request all partial rewards at the last timestep and return them as
      # a list.
      last_rewards = episode_1.get_rewards(
          partial=True, consider_buffer=False, as_list=True
      )
      self.assertIn("agent_9", last_rewards[0])
      self.assertEqual(
          episode_1.partial_rewards["agent_9"][-1],
          last_rewards[0]["agent_9"],
      )
      self.assertIn("agent_0", last_rewards[0])
      self.assertEqual(
          episode_1.partial_rewards["agent_0"][-1],
          last_rewards[0]["agent_0"],
      )
      self.assertIn("agent_2", last_rewards[0])
      self.assertEqual(
          episode_1.partial_rewards["agent_2"][-1],
          last_rewards[0]["agent_2"],
      )
      self.assertIn("agent_8", last_rewards[0])
      self.assertEqual(
          episode_1.partial_rewards["agent_8"][-1],
          last_rewards[0]["agent_8"],
      )
      self.assertIn("agent_4", last_rewards[0])
      self.assertEqual(
          episode_1.partial_rewards["agent_4"][-1],
          last_rewards[0]["agent_4"],
      )
      self.assertIn("agent_3", last_rewards[0])
      self.assertEqual(
          episode_1.partial_rewards["agent_3"][-1],
          last_rewards[0]["agent_3"],
      )
      self.assertIn("agent_6", last_rewards[0])
      self.assertEqual(
          episode_1.partial_rewards["agent_6"][-1],
          last_rewards[0]["agent_6"],
      )
      self.assertIn("agent_7", last_rewards[0])
      self.assertEqual(
          episode_1.partial_rewards["agent_7"][-1],
          last_rewards[0]["agent_7"],
      )
      # Assert that all the other agents are not in the returned rewards.
      self.assertNotIn("agent_1", last_rewards[0])
  
      # Request the last two indices, but consider only partial rewards.
      last_rewards = episode_1.get_rewards(
          [-1, -2], partial=True, consider_buffer=False
      )
      self.assertIn("agent_9", last_rewards)
      self.assertListEqual(
          episode_1.partial_rewards_t["agent_9"][-2:], [99, 100]
      )
      self.assertEqual(
          episode_1.partial_rewards["agent_9"][-1:-3:-1],
          last_rewards["agent_9"],
      )
      self.assertIn("agent_0", last_rewards)
      self.assertListEqual(
          episode_1.partial_rewards_t["agent_0"][-2:], [99, 100]
      )
      self.assertEqual(
          episode_1.partial_rewards["agent_0"][-1:-3:-1],
          last_rewards["agent_0"],
      )
      self.assertIn("agent_2", last_rewards)
      self.assertListEqual(
          episode_1.partial_rewards_t["agent_2"][-2:], [99, 100]
      )
      self.assertEqual(
          episode_1.partial_rewards["agent_2"][-1:-3:-1],
          last_rewards["agent_2"],
      )
      self.assertIn("agent_8", last_rewards)
      self.assertListEqual(
          episode_1.partial_rewards_t["agent_8"][-2:], [99, 100]
      )
      self.assertEqual(
          episode_1.partial_rewards["agent_8"][-1:-3:-1],
          last_rewards["agent_8"],
      )
      self.assertIn("agent_4", last_rewards)
      self.assertListEqual(
          episode_1.partial_rewards_t["agent_4"][-2:], [99, 100]
      )
      self.assertEqual(
          episode_1.partial_rewards["agent_4"][-1:-3:-1],
          last_rewards["agent_4"],
      )
      self.assertIn("agent_3", last_rewards)
      self.assertListEqual(
          episode_1.partial_rewards_t["agent_3"][-2:], [99, 100]
      )
      self.assertEqual(
          episode_1.partial_rewards["agent_3"][-1:-3:-1],
          last_rewards["agent_3"],
      )
      self.assertIn("agent_6", last_rewards)
      self.assertListEqual(
          episode_1.partial_rewards_t["agent_6"][-2:], [99, 100]
      )
      self.assertEqual(
          episode_1.partial_rewards["agent_6"][-1:-3:-1],
          last_rewards["agent_6"],
      )
      self.assertIn("agent_7", last_rewards)
      self.assertListEqual(
          episode_1.partial_rewards_t["agent_7"][-2:], [99, 100]
      )
      self.assertEqual(
          episode_1.partial_rewards["agent_7"][-1:-3:-1],
          last_rewards["agent_7"],
      )
      # Assert that all the other agents are not in the returned rewards.
      self.assertNotIn("agent_1", last_rewards)
  
      # At last, request the last two indices for only partial rewards and return
      # them as list.
      last_rewards = episode_1.get_rewards(
          [-1, -2], partial=True, consider_buffer=False, as_list=True
      )
      self.assertIn("agent_9", last_rewards[0])
      self.assertIn("agent_9", last_rewards[1])
      self.assertEqual(
          episode_1.partial_rewards["agent_9"][-1],
          last_rewards[0]["agent_9"],
      )
      self.assertEqual(
          episode_1.partial_rewards["agent_9"][-2],
          last_rewards[1]["agent_9"],
      )
      self.assertIn("agent_0", last_rewards[0])
      self.assertIn("agent_0", last_rewards[1])
      self.assertEqual(
          episode_1.partial_rewards["agent_0"][-1],
          last_rewards[0]["agent_0"],
      )
      self.assertEqual(
          episode_1.partial_rewards["agent_0"][-2],
          last_rewards[1]["agent_0"],
      )
      self.assertIn("agent_2", last_rewards[0])
      self.assertIn("agent_2", last_rewards[1])
      self.assertEqual(
          episode_1.partial_rewards["agent_2"][-1],
          last_rewards[0]["agent_2"],
      )
      self.assertEqual(
          episode_1.partial_rewards["agent_2"][-2],
          last_rewards[1]["agent_2"],
      )
      self.assertIn("agent_8", last_rewards[0])
      self.assertIn("agent_8", last_rewards[1])
      self.assertEqual(
          episode_1.partial_rewards["agent_8"][-1],
          last_rewards[0]["agent_8"],
      )
      self.assertEqual(
          episode_1.partial_rewards["agent_8"][-2],
          last_rewards[1]["agent_8"],
      )
      self.assertIn("agent_4", last_rewards[0])
      self.assertIn("agent_4", last_rewards[1])
      self.assertEqual(
          episode_1.partial_rewards["agent_4"][-1],
          last_rewards[0]["agent_4"],
      )
      self.assertEqual(
          episode_1.partial_rewards["agent_4"][-2],
          last_rewards[1]["agent_4"],
      )
      self.assertIn("agent_3", last_rewards[0])
      self.assertIn("agent_3", last_rewards[1])
      self.assertEqual(
          episode_1.partial_rewards["agent_3"][-1],
          last_rewards[0]["agent_3"],
      )
      self.assertEqual(
          episode_1.partial_rewards["agent_3"][-2],
          last_rewards[1]["agent_3"],
      )
      self.assertIn("agent_6", last_rewards[0])
      self.assertIn("agent_6", last_rewards[1])
      self.assertEqual(
          episode_1.partial_rewards["agent_6"][-1],
          last_rewards[0]["agent_6"],
      )
      self.assertEqual(
          episode_1.partial_rewards["agent_6"][-2],
          last_rewards[1]["agent_6"],
      )
      self.assertIn("agent_7", last_rewards[0])
      self.assertIn("agent_7", last_rewards[1])
      self.assertEqual(
          episode_1.partial_rewards["agent_7"][-1],
          last_rewards[0]["agent_7"],
      )
      self.assertEqual(
          episode_1.partial_rewards["agent_7"][-2],
          last_rewards[1]["agent_7"],
      )
      # Assert that all the other agents are not in the returned rewards.
      self.assertNotIn("agent_1", last_rewards[0])
      self.assertNotIn("agent_1", last_rewards[1])
  
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 7 for calling the TestMultiAgentEpisode obj's underlying member method code with example usage: obj.test_getters() without return types