code: def test_returning_model_based_rollouts_data(self):
      # TODO(avnishn): This test only works with the old api
  
      class ModelBasedPolicy(DQNTFPolicy):
          def compute_actions_from_input_dict(
              self,
              input_dict,
              explore=None,
              timestep=None,
              episodes=None,
              **kwargs,
          ):
              obs_batch = input_dict["obs"]
              # In policy loss initialization phase, no episodes are passed
              # in.
              if episodes is not None:
                  # Pretend we did a model-based rollout and want to return
                  # the extra trajectory.
                  env_id = episodes[0].env_id
                  fake_eps = Episode(
                      episodes[0].policy_map,
                      episodes[0].policy_mapping_fn,
                      lambda: None,
                      lambda x: None,
                      env_id,
                  )
                  builder = get_global_worker().sampler.sample_collector
                  agent_id = "extra_0"
                  policy_id = "p1"  # use p1 so we can easily check it
                  builder.add_init_obs(
                      episode=fake_eps,
                      agent_id=agent_id,
                      policy_id=policy_id,
                      env_id=env_id,
                      init_obs=obs_batch[0],
                      init_infos={},
                  )
                  for t in range(4):
                      builder.add_action_reward_next_obs(
                          episode_id=fake_eps.episode_id,
                          agent_id=agent_id,
                          env_id=env_id,
                          policy_id=policy_id,
                          agent_done=t == 3,
                          values=dict(
                              t=t,
                              actions=0,
                              rewards=0,
                              terminateds=False,
                              truncateds=t == 3,
                              infos={},
                              new_obs=obs_batch[0],
                          ),
                      )
                  batch = builder.postprocess_episode(
                      episode=fake_eps, build=True
                  )
                  episodes[0].add_extra_batch(batch)
  
              # Just return zeros for actions
              return [0] * len(obs_batch), [], {}
  
      ev = RolloutWorker(
          env_creator=lambda _: MultiAgentCartPole({"num_agents": 2}),
          default_policy_class=ModelBasedPolicy,
          config=DQNConfig()
          .framework("tf")
          .rollouts(
              rollout_fragment_length=5,
              num_rollout_workers=0,
              enable_connectors=False,  # only works with old episode API
          )
          .multi_agent(
              policies={"p0", "p1"},
              policy_mapping_fn=lambda agent_id, episode, worker, **kwargs: "p0",
          ),
      )
      batch = ev.sample()
      # 5 environment steps (rollout_fragment_length).
      self.assertEqual(batch.count, 5)
      # 10 agent steps for p0: 2 agents, both using p0 as their policy.
      self.assertEqual(batch.policy_batches["p0"].count, 10)
      # 20 agent steps for p1: Each time both(!) agents takes 1 step,
      # p1 takes 4: 5 (rollout-fragment length) * 4 = 20
      self.assertEqual(batch.policy_batches["p1"].count, 20)
  
imports: import gymnasium
  import numpy
  import random
  import tree
  import unittest
  import ray
  from ray.tune.registry import register_env
  from ray.rllib.algorithms.algorithm_config import AlgorithmConfig
  from ray.rllib.algorithms.dqn.dqn import DQNConfig
  from ray.rllib.algorithms.dqn.dqn_tf_policy import DQNTFPolicy
  from ray.rllib.algorithms.ppo import PPOConfig
  from ray.rllib.env.multi_agent_env import make_multi_agent, MultiAgentEnvWrapper
  from ray.rllib.evaluation.episode import Episode
  from ray.rllib.evaluation.rollout_worker import get_global_worker, RolloutWorker
  from ray.rllib.evaluation.tests.test_rollout_worker import MockPolicy
  from ray.rllib.examples.policy.random_policy import RandomPolicy
  from ray.rllib.examples.env.multi_agent import MultiAgentCartPole, BasicMultiAgent, EarlyDoneMultiAgent, FlexAgentsMultiAgent, RoundRobinMultiAgent, SometimesZeroAgentsMultiAgent
  from ray.rllib.policy.policy import PolicySpec
  from ray.rllib.policy.sample_batch import convert_ma_batch_to_sample_batch
  from ray.rllib.tests.test_nested_observation_spaces import NestedMultiAgentEnv
  from ray.rllib.utils.numpy import one_hot
  from ray.rllib.utils.test_utils import check
label_desc: Use image node_id 16 for calling the TestMultiAgentEnv obj's underlying member method code with example usage: obj.test_returning_model_based_rollouts_data() and returns: unknown,  list,  dict