code: def test_add_env_step(self):
      # Create an environment and add the initial observations, infos, and states.
      env = MultiAgentTestEnv()
      episode = MultiAgentEpisode(agent_ids=env.get_agent_ids())
  
      obs, infos = env.reset(seed=0)
      episode.add_env_reset(
          observations=obs,
          infos=infos,
      )
  
      # Sample 100 timesteps and add them to the episode.
      for i in range(100):
          action = {
              agent_id: i + 1
              for agent_id in obs
              if agent_id in env._agents_alive
          }
          obs, reward, terminated, truncated, info = env.step(action)
  
          episode.add_env_step(
              observations=obs,
              actions=action,
              rewards=reward,
              infos=info,
              terminateds=terminated,
              truncateds=truncated,
              extra_model_outputs={
                  agent_id: {"extra": 10.5} for agent_id in action
              },
          )
  
      # Assert that the timestep is at 100.
      self.assertEqual(episode.t, 100)
      # Ensure that the episode is not done yet.
      self.assertFalse(episode.is_done)
      # Ensure that agent 1 and agent 5 are indeed done.
      self.assertTrue(episode.agent_episodes["agent_1"].is_done)
      self.assertTrue(episode.agent_episodes["agent_5"].is_done)
      # Also ensure that their buffers are all empty:
      for agent_id in ["agent_1", "agent_5"]:
          self.assertTrue(
              episode.agent_buffers[agent_id]["actions"].empty()
          )
          self.assertTrue(
              episode.agent_buffers[agent_id]["rewards"].empty()
          )
          self.assertTrue(
              episode.agent_buffers[agent_id][
                  "extra_model_outputs"
              ].empty()
          )
      # Ensure that the maximum timestep in the global timestep mapping
      # is 100.
      highest_timestep = max(
          [
              max(timesteps)
              for timesteps in episode.global_t_to_local_t.values()
          ]
      )
      self.assertGreaterEqual(100, highest_timestep)
  
      # Run another 100 timesteps.
      for i in range(100, 200):
          action = {
              agent_id: i + 1
              for agent_id in obs
              if agent_id in env._agents_alive
          }
          obs, reward, terminated, truncated, info = env.step(action)
          episode.add_env_step(
              observations=obs,
              actions=action,
              rewards=reward,
              infos=info,
              terminateds=terminated,
              truncateds=truncated,
              extra_model_outputs={
                  agent_id: {"extra": 10.5} for agent_id in action
              },
          )
  
      # Assert that the environment is done.
      self.assertTrue(truncated["__all__"])
      # Assert that each agent is done.
      for agent_id in episode._agent_ids:
          self.assertTrue(episode.agent_episodes[agent_id].is_done)
      # Assert that agent 1 and agent 5 have no observations/actions/etc.
      # after the timesteps in which they terminated.
      self.assertGreaterEqual(
          50, episode.agent_episodes["agent_1"].observations[-1]
      )
      self.assertGreaterEqual(
          50, episode.agent_episodes["agent_1"].actions[-1]
      )
      self.assertGreaterEqual(
          100, episode.agent_episodes["agent_5"].observations[-1]
      )
      self.assertGreaterEqual(
          100, episode.agent_episodes["agent_5"].actions[-1]
      )
  
      # Now test, if agents that have never stepped are handled correctly.
      # agent 5 will be the agent that never stepped.
      (
          observations,
          actions,
          rewards,
          terminated,
          truncated,
          infos,
      ) = self._mock_multi_agent_records()
  
      episode = MultiAgentEpisode(
          agent_ids=[
              "agent_1",
              "agent_2",
              "agent_3",
              "agent_4",
              "agent_5",
          ],
          observations=observations,
          actions=actions,
          rewards=rewards,
          infos=infos,
          terminateds=terminated,
          truncateds=truncated,
      )
      # Now test that intermediate rewards will get recorded and actions buffered.
      action = {"agent_2": 3, "agent_4": 3}
      observation = {"agent_1": 3, "agent_2": 3}
      reward = {
          "agent_1": 1.0,
          "agent_2": 1.0,
          "agent_3": 1.0,
          "agent_5": 1.0,
      }
      infos = {"agent_1": {}, "agent_2": {}}
      terminated = {k: False for k in observation.keys()}
      terminated.update({"__all__": False})
      truncated = {k: False for k in observation.keys()}
      truncated.update({"__all__": False})
      episode.add_env_step(
          observations=observation,
          actions=action,
          rewards=reward,
          infos=infos,
          terminateds=terminated,
          truncateds=truncated,
      )
      # Assert that the action buffer for agent 4 is full.
      # Note, agent 4 acts, but receives no observation.
      # Note also, all other buffers are always full, due to their defaults.
      self.assertTrue(
          episode.agent_buffers["agent_4"]["actions"].full()
      )
      # Assert that the reward buffers of agents 3 and 5 are at 1.0.
      self.assertEquals(
          episode.agent_buffers["agent_3"]["rewards"].get_nowait(), 1.0
      )
      self.assertEquals(
          episode.agent_buffers["agent_5"]["rewards"].get_nowait(), 1.0
      )
      # Now, refill the buffers.
      episode.agent_buffers["agent_3"]["rewards"].put_nowait(1.0)
      episode.agent_buffers["agent_5"]["rewards"].put_nowait(1.0)
  
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 5 for calling the TestMultiAgentEpisode obj's underlying member method code with example usage: obj.test_add_env_step() without return types