prompt: We are currently solving the following issue within our repository. Here is the issue text:
  --- BEGIN ISSUE ---
  [rllib] make rollout script support multiagent 
  Hi,
  
  If I'm correct, only single agent/policy is currently supported in rollout.py. For instance https://github.com/ray-project/ray/blob/2e30f7ba386e716bf80f019dcd473b67d83abb95/python/ray/rllib/rollout.py#L109-L110 references default policy to check if policy uses lstm, which fails when a multi agent configuration is loaded.
  
  Thanks!
  
  --- END ISSUE ---
  
  Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
  --- BEGIN FILES ---
  Path: `python/ray/rllib/rollout.py`
  Content:
  ```
  1 #!/usr/bin/env python
  2 
  3 from __future__ import absolute_import
  4 from __future__ import division
  5 from __future__ import print_function
  6 
  7 import argparse
  8 import json
  9 import os
  10 import pickle
  11 
  12 import gym
  13 import ray
  14 from ray.rllib.agents.registry import get_agent_class
  15 
  16 EXAMPLE_USAGE = """
  17 Example Usage via RLlib CLI:
  18     rllib rollout /tmp/ray/checkpoint_dir/checkpoint-0 --run DQN
  19     --env CartPole-v0 --steps 1000000 --out rollouts.pkl
  20 
  21 Example Usage via executable:
  22     ./rollout.py /tmp/ray/checkpoint_dir/checkpoint-0 --run DQN
  23     --env CartPole-v0 --steps 1000000 --out rollouts.pkl
  24 """
  25 
  26 # Note: if you use any custom models or envs, register them here first, e.g.:
  27 #
  28 # ModelCatalog.register_custom_model("pa_model", ParametricActionsModel)
  29 # register_env("pa_cartpole", lambda _: ParametricActionCartpole(10))
  30 
  31 
  32 def create_parser(parser_creator=None):
  33     parser_creator = parser_creator or argparse.ArgumentParser
  34     parser = parser_creator(
  35         formatter_class=argparse.RawDescriptionHelpFormatter,
  36         description="Roll out a reinforcement learning agent "
  37         "given a checkpoint.",
  38         epilog=EXAMPLE_USAGE)
  39 
  40     parser.add_argument(
  41         "checkpoint", type=str, help="Checkpoint from which to roll out.")
  42     required_named = parser.add_argument_group("required named arguments")
  43     required_named.add_argument(
  44         "--run",
  45         type=str,
  46         required=True,
  47         help="The algorithm or model to train. This may refer to the name "
  48         "of a built-on algorithm (e.g. RLLib's DQN or PPO), or a "
  49         "user-defined trainable function or class registered in the "
  50         "tune registry.")
  51     required_named.add_argument(
  52         "--env", type=str, help="The gym environment to use.")
  53     parser.add_argument(
  54         "--no-render",
  55         default=False,
  56         action="store_const",
  57         const=True,
  58         help="Surpress rendering of the environment.")
  59     parser.add_argument(
  60         "--steps", default=10000, help="Number of steps to roll out.")
  61     parser.add_argument("--out", default=None, help="Output filename.")
  62     parser.add_argument(
  63         "--config",
  64         default="{}",
  65         type=json.loads,
  66         help="Algorithm-specific configuration (e.g. env, hyperparams). "
  67         "Surpresses loading of configuration from checkpoint.")
  68     return parser
  69 
  70 
  71 def run(args, parser):
  72     config = args.config
  73     if not config:
  74         # Load configuration from file
  75         config_dir = os.path.dirname(args.checkpoint)
  76         config_path = os.path.join(config_dir, "params.json")
  77         if not os.path.exists(config_path):
  78             config_path = os.path.join(config_dir, "../params.json")
  79         if not os.path.exists(config_path):
  80             raise ValueError(
  81                 "Could not find params.json in either the checkpoint dir or "
  82                 "its parent directory.")
  83         with open(config_path) as f:
  84             config = json.load(f)
  85         if "num_workers" in config:
  86             config["num_workers"] = min(2, config["num_workers"])
  87 
  88     if not args.env:
  89         if not config.get("env"):
  90             parser.error("the following arguments are required: --env")
  91         args.env = config.get("env")
  92 
  93     ray.init()
  94 
  95     cls = get_agent_class(args.run)
  96     agent = cls(env=args.env, config=config)
  97     agent.restore(args.checkpoint)
  98     num_steps = int(args.steps)
  99     rollout(agent, args.env, num_steps, args.out, args.no_render)
  100 
  101 
  102 def rollout(agent, env_name, num_steps, out=None, no_render=True):
  103     if hasattr(agent, "local_evaluator"):
  104         env = agent.local_evaluator.env
  105     else:
  106         env = gym.make(env_name)
  107 
  108     if hasattr(agent, "local_evaluator"):
  109         state_init = agent.local_evaluator.policy_map[
  110             "default"].get_initial_state()
  111     else:
  112         state_init = []
  113     if state_init:
  114         use_lstm = True
  115     else:
  116         use_lstm = False
  117 
  118     if out is not None:
  119         rollouts = []
  120     steps = 0
  121     while steps < (num_steps or steps + 1):
  122         if out is not None:
  123             rollout = []
  124         state = env.reset()
  125         done = False
  126         reward_total = 0.0
  127         while not done and steps < (num_steps or steps + 1):
  128             if use_lstm:
  129                 action, state_init, logits = agent.compute_action(
  130                     state, state=state_init)
  131             else:
  132                 action = agent.compute_action(state)
  133             next_state, reward, done, _ = env.step(action)
  134             reward_total += reward
  135             if not no_render:
  136                 env.render()
  137             if out is not None:
  138                 rollout.append([state, action, next_state, reward, done])
  139             steps += 1
  140             state = next_state
  141         if out is not None:
  142             rollouts.append(rollout)
  143         print("Episode reward", reward_total)
  144     if out is not None:
  145         pickle.dump(rollouts, open(out, "wb"))
  146 
  147 
  148 if __name__ == "__main__":
  149     parser = create_parser()
  150     args = parser.parse_args()
  151     run(args, parser)
  152
  ```
  --- END FILES ---
  
  Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.
  
  Here is an example:
  
  ```diff
  diff --git a/examples/server_async.py b/examples/server_async.py
  --- a/examples/server_async.py
  +++ b/examples/server_async.py
  @@ -313,4 +313,4 @@
   
   
   if __name__ == "__main__":
  -    asyncio.run(run_async_server("."), debug=True)
  +    asyncio.run(run_async_server(), debug=True)
  diff --git a/examples/server_sync.py b/examples/server_sync.py
  --- a/examples/server_sync.py
  +++ b/examples/server_sync.py
  @@ -313,5 +313,5 @@
   
   
   if __name__ == "__main__":
  -    server = run_sync_server(".")
  +    server = run_sync_server()
       server.shutdown()
  
  ```
  
golden_diff: diff --git a/python/ray/rllib/rollout.py b/python/ray/rllib/rollout.py
  --- a/python/ray/rllib/rollout.py
  +++ b/python/ray/rllib/rollout.py
  @@ -73,15 +73,15 @@
       if not config:
           # Load configuration from file
           config_dir = os.path.dirname(args.checkpoint)
  -        config_path = os.path.join(config_dir, "params.json")
  +        config_path = os.path.join(config_dir, "params.pkl")
           if not os.path.exists(config_path):
  -            config_path = os.path.join(config_dir, "../params.json")
  +            config_path = os.path.join(config_dir, "../params.pkl")
           if not os.path.exists(config_path):
               raise ValueError(
  -                "Could not find params.json in either the checkpoint dir or "
  +                "Could not find params.pkl in either the checkpoint dir or "
                   "its parent directory.")
  -        with open(config_path) as f:
  -            config = json.load(f)
  +        with open(config_path, 'rb') as f:
  +            config = pickle.load(f)
           if "num_workers" in config:
               config["num_workers"] = min(2, config["num_workers"])
   
  @@ -102,18 +102,18 @@
   def rollout(agent, env_name, num_steps, out=None, no_render=True):
       if hasattr(agent, "local_evaluator"):
           env = agent.local_evaluator.env
  +        multiagent = agent.local_evaluator.multiagent
  +        if multiagent:
  +            policy_agent_mapping = agent.config["multiagent"][
  +                "policy_mapping_fn"]
  +            mapping_cache = {}
  +        policy_map = agent.local_evaluator.policy_map
  +        state_init = {p: m.get_initial_state() for p, m in policy_map.items()}
  +        use_lstm = {p: len(s) > 0 for p, s in state_init.items()}
       else:
           env = gym.make(env_name)
  -
  -    if hasattr(agent, "local_evaluator"):
  -        state_init = agent.local_evaluator.policy_map[
  -            "default"].get_initial_state()
  -    else:
  -        state_init = []
  -    if state_init:
  -        use_lstm = True
  -    else:
  -        use_lstm = False
  +        multiagent = False
  +        use_lstm = {'default': False}
   
       if out is not None:
           rollouts = []
  @@ -125,13 +125,39 @@
           done = False
           reward_total = 0.0
           while not done and steps < (num_steps or steps + 1):
  -            if use_lstm:
  -                action, state_init, logits = agent.compute_action(
  -                    state, state=state_init)
  +            if multiagent:
  +                action_dict = {}
  +                for agent_id in state.keys():
  +                    a_state = state[agent_id]
  +                    if a_state is not None:
  +                        policy_id = mapping_cache.setdefault(
  +                            agent_id, policy_agent_mapping(agent_id))
  +                        p_use_lstm = use_lstm[policy_id]
  +                        if p_use_lstm:
  +                            a_action, p_state_init, _ = agent.compute_action(
  +                                a_state,
  +                                state=state_init[policy_id],
  +                                policy_id=policy_id)
  +                            state_init[policy_id] = p_state_init
  +                        else:
  +                            a_action = agent.compute_action(
  +                                a_state, policy_id=policy_id)
  +                        action_dict[agent_id] = a_action
  +                action = action_dict
               else:
  -                action = agent.compute_action(state)
  +                if use_lstm["default"]:
  +                    action, state_init, _ = agent.compute_action(
  +                        state, state=state_init)
  +                else:
  +                    action = agent.compute_action(state)
  +
               next_state, reward, done, _ = env.step(action)
  -            reward_total += reward
  +
  +            if multiagent:
  +                done = done["__all__"]
  +                reward_total += sum(reward.values())
  +            else:
  +                reward_total += reward
               if not no_render:
                   env.render()
               if out is not None:
  @@ -141,6 +167,7 @@
           if out is not None:
               rollouts.append(rollout)
           print("Episode reward", reward_total)
  +
       if out is not None:
           pickle.dump(rollouts, open(out, "wb"))
  
verification_info: {"golden_diff": "diff --git a/python/ray/rllib/rollout.py b/python/ray/rllib/rollout.py\n--- a/python/ray/rllib/rollout.py\n+++ b/python/ray/rllib/rollout.py\n@@ -73,15 +73,15 @@\n     if not config:\n         # Load configuration from file\n         config_dir = os.path.dirname(args.checkpoint)\n-        config_path = os.path.join(config_dir, \"params.json\")\n+        config_path = os.path.join(config_dir, \"params.pkl\")\n         if not os.path.exists(config_path):\n-            config_path = os.path.join(config_dir, \"../params.json\")\n+            config_path = os.path.join(config_dir, \"../params.pkl\")\n         if not os.path.exists(config_path):\n             raise ValueError(\n-                \"Could not find params.json in either the checkpoint dir or \"\n+                \"Could not find params.pkl in either the checkpoint dir or \"\n                 \"its parent directory.\")\n-        with open(config_path) as f:\n-            config = json.load(f)\n+        with open(config_path, 'rb') as f:\n+            config = pickle.load(f)\n         if \"num_workers\" in config:\n             config[\"num_workers\"] = min(2, config[\"num_workers\"])\n \n@@ -102,18 +102,18 @@\n def rollout(agent, env_name, num_steps, out=None, no_render=True):\n     if hasattr(agent, \"local_evaluator\"):\n         env = agent.local_evaluator.env\n+        multiagent = agent.local_evaluator.multiagent\n+        if multiagent:\n+            policy_agent_mapping = agent.config[\"multiagent\"][\n+                \"policy_mapping_fn\"]\n+            mapping_cache = {}\n+        policy_map = agent.local_evaluator.policy_map\n+        state_init = {p: m.get_initial_state() for p, m in policy_map.items()}\n+        use_lstm = {p: len(s) > 0 for p, s in state_init.items()}\n     else:\n         env = gym.make(env_name)\n-\n-    if hasattr(agent, \"local_evaluator\"):\n-        state_init = agent.local_evaluator.policy_map[\n-            \"default\"].get_initial_state()\n-    else:\n-        state_init = []\n-    if state_init:\n-        use_lstm = True\n-    else:\n-        use_lstm = False\n+        multiagent = False\n+        use_lstm = {'default': False}\n \n     if out is not None:\n         rollouts = []\n@@ -125,13 +125,39 @@\n         done = False\n         reward_total = 0.0\n         while not done and steps < (num_steps or steps + 1):\n-            if use_lstm:\n-                action, state_init, logits = agent.compute_action(\n-                    state, state=state_init)\n+            if multiagent:\n+                action_dict = {}\n+                for agent_id in state.keys():\n+                    a_state = state[agent_id]\n+                    if a_state is not None:\n+                        policy_id = mapping_cache.setdefault(\n+                            agent_id, policy_agent_mapping(agent_id))\n+                        p_use_lstm = use_lstm[policy_id]\n+                        if p_use_lstm:\n+                            a_action, p_state_init, _ = agent.compute_action(\n+                                a_state,\n+                                state=state_init[policy_id],\n+                                policy_id=policy_id)\n+                            state_init[policy_id] = p_state_init\n+                        else:\n+                            a_action = agent.compute_action(\n+                                a_state, policy_id=policy_id)\n+                        action_dict[agent_id] = a_action\n+                action = action_dict\n             else:\n-                action = agent.compute_action(state)\n+                if use_lstm[\"default\"]:\n+                    action, state_init, _ = agent.compute_action(\n+                        state, state=state_init)\n+                else:\n+                    action = agent.compute_action(state)\n+\n             next_state, reward, done, _ = env.step(action)\n-            reward_total += reward\n+\n+            if multiagent:\n+                done = done[\"__all__\"]\n+                reward_total += sum(reward.values())\n+            else:\n+                reward_total += reward\n             if not no_render:\n                 env.render()\n             if out is not None:\n@@ -141,6 +167,7 @@\n         if out is not None:\n             rollouts.append(rollout)\n         print(\"Episode reward\", reward_total)\n+\n     if out is not None:\n         pickle.dump(rollouts, open(out, \"wb\"))\n", "issue": "[rllib] make rollout script support multiagent \nHi,\r\n\r\nIf I'm correct, only single agent/policy is currently supported in rollout.py. For instance https://github.com/ray-project/ray/blob/2e30f7ba386e716bf80f019dcd473b67d83abb95/python/ray/rllib/rollout.py#L109-L110 references default policy to check if policy uses lstm, which fails when a multi agent configuration is loaded.\r\n\r\nThanks!\n", "before_files": [{"content": "#!/usr/bin/env python\n\nfrom __future__ import absolute_import\nfrom __future__ import division\nfrom __future__ import print_function\n\nimport argparse\nimport json\nimport os\nimport pickle\n\nimport gym\nimport ray\nfrom ray.rllib.agents.registry import get_agent_class\n\nEXAMPLE_USAGE = \"\"\"\nExample Usage via RLlib CLI:\n    rllib rollout /tmp/ray/checkpoint_dir/checkpoint-0 --run DQN\n    --env CartPole-v0 --steps 1000000 --out rollouts.pkl\n\nExample Usage via executable:\n    ./rollout.py /tmp/ray/checkpoint_dir/checkpoint-0 --run DQN\n    --env CartPole-v0 --steps 1000000 --out rollouts.pkl\n\"\"\"\n\n# Note: if you use any custom models or envs, register them here first, e.g.:\n#\n# ModelCatalog.register_custom_model(\"pa_model\", ParametricActionsModel)\n# register_env(\"pa_cartpole\", lambda _: ParametricActionCartpole(10))\n\n\ndef create_parser(parser_creator=None):\n    parser_creator = parser_creator or argparse.ArgumentParser\n    parser = parser_creator(\n        formatter_class=argparse.RawDescriptionHelpFormatter,\n        description=\"Roll out a reinforcement learning agent \"\n        \"given a checkpoint.\",\n        epilog=EXAMPLE_USAGE)\n\n    parser.add_argument(\n        \"checkpoint\", type=str, help=\"Checkpoint from which to roll out.\")\n    required_named = parser.add_argument_group(\"required named arguments\")\n    required_named.add_argument(\n        \"--run\",\n        type=str,\n        required=True,\n        help=\"The algorithm or model to train. This may refer to the name \"\n        \"of a built-on algorithm (e.g. RLLib's DQN or PPO), or a \"\n        \"user-defined trainable function or class registered in the \"\n        \"tune registry.\")\n    required_named.add_argument(\n        \"--env\", type=str, help=\"The gym environment to use.\")\n    parser.add_argument(\n        \"--no-render\",\n        default=False,\n        action=\"store_const\",\n        const=True,\n        help=\"Surpress rendering of the environment.\")\n    parser.add_argument(\n        \"--steps\", default=10000, help=\"Number of steps to roll out.\")\n    parser.add_argument(\"--out\", default=None, help=\"Output filename.\")\n    parser.add_argument(\n        \"--config\",\n        default=\"{}\",\n        type=json.loads,\n        help=\"Algorithm-specific configuration (e.g. env, hyperparams). \"\n        \"Surpresses loading of configuration from checkpoint.\")\n    return parser\n\n\ndef run(args, parser):\n    config = args.config\n    if not config:\n        # Load configuration from file\n        config_dir = os.path.dirname(args.checkpoint)\n        config_path = os.path.join(config_dir, \"params.json\")\n        if not os.path.exists(config_path):\n            config_path = os.path.join(config_dir, \"../params.json\")\n        if not os.path.exists(config_path):\n            raise ValueError(\n                \"Could not find params.json in either the checkpoint dir or \"\n                \"its parent directory.\")\n        with open(config_path) as f:\n            config = json.load(f)\n        if \"num_workers\" in config:\n            config[\"num_workers\"] = min(2, config[\"num_workers\"])\n\n    if not args.env:\n        if not config.get(\"env\"):\n            parser.error(\"the following arguments are required: --env\")\n        args.env = config.get(\"env\")\n\n    ray.init()\n\n    cls = get_agent_class(args.run)\n    agent = cls(env=args.env, config=config)\n    agent.restore(args.checkpoint)\n    num_steps = int(args.steps)\n    rollout(agent, args.env, num_steps, args.out, args.no_render)\n\n\ndef rollout(agent, env_name, num_steps, out=None, no_render=True):\n    if hasattr(agent, \"local_evaluator\"):\n        env = agent.local_evaluator.env\n    else:\n        env = gym.make(env_name)\n\n    if hasattr(agent, \"local_evaluator\"):\n        state_init = agent.local_evaluator.policy_map[\n            \"default\"].get_initial_state()\n    else:\n        state_init = []\n    if state_init:\n        use_lstm = True\n    else:\n        use_lstm = False\n\n    if out is not None:\n        rollouts = []\n    steps = 0\n    while steps < (num_steps or steps + 1):\n        if out is not None:\n            rollout = []\n        state = env.reset()\n        done = False\n        reward_total = 0.0\n        while not done and steps < (num_steps or steps + 1):\n            if use_lstm:\n                action, state_init, logits = agent.compute_action(\n                    state, state=state_init)\n            else:\n                action = agent.compute_action(state)\n            next_state, reward, done, _ = env.step(action)\n            reward_total += reward\n            if not no_render:\n                env.render()\n            if out is not None:\n                rollout.append([state, action, next_state, reward, done])\n            steps += 1\n            state = next_state\n        if out is not None:\n            rollouts.append(rollout)\n        print(\"Episode reward\", reward_total)\n    if out is not None:\n        pickle.dump(rollouts, open(out, \"wb\"))\n\n\nif __name__ == \"__main__\":\n    parser = create_parser()\n    args = parser.parse_args()\n    run(args, parser)\n", "path": "python/ray/rllib/rollout.py"}], "after_files": [{"content": "#!/usr/bin/env python\n\nfrom __future__ import absolute_import\nfrom __future__ import division\nfrom __future__ import print_function\n\nimport argparse\nimport json\nimport os\nimport pickle\n\nimport gym\nimport ray\nfrom ray.rllib.agents.registry import get_agent_class\n\nEXAMPLE_USAGE = \"\"\"\nExample Usage via RLlib CLI:\n    rllib rollout /tmp/ray/checkpoint_dir/checkpoint-0 --run DQN\n    --env CartPole-v0 --steps 1000000 --out rollouts.pkl\n\nExample Usage via executable:\n    ./rollout.py /tmp/ray/checkpoint_dir/checkpoint-0 --run DQN\n    --env CartPole-v0 --steps 1000000 --out rollouts.pkl\n\"\"\"\n\n# Note: if you use any custom models or envs, register them here first, e.g.:\n#\n# ModelCatalog.register_custom_model(\"pa_model\", ParametricActionsModel)\n# register_env(\"pa_cartpole\", lambda _: ParametricActionCartpole(10))\n\n\ndef create_parser(parser_creator=None):\n    parser_creator = parser_creator or argparse.ArgumentParser\n    parser = parser_creator(\n        formatter_class=argparse.RawDescriptionHelpFormatter,\n        description=\"Roll out a reinforcement learning agent \"\n        \"given a checkpoint.\",\n        epilog=EXAMPLE_USAGE)\n\n    parser.add_argument(\n        \"checkpoint\", type=str, help=\"Checkpoint from which to roll out.\")\n    required_named = parser.add_argument_group(\"required named arguments\")\n    required_named.add_argument(\n        \"--run\",\n        type=str,\n        required=True,\n        help=\"The algorithm or model to train. This may refer to the name \"\n        \"of a built-on algorithm (e.g. RLLib's DQN or PPO), or a \"\n        \"user-defined trainable function or class registered in the \"\n        \"tune registry.\")\n    required_named.add_argument(\n        \"--env\", type=str, help=\"The gym environment to use.\")\n    parser.add_argument(\n        \"--no-render\",\n        default=False,\n        action=\"store_const\",\n        const=True,\n        help=\"Surpress rendering of the environment.\")\n    parser.add_argument(\n        \"--steps\", default=10000, help=\"Number of steps to roll out.\")\n    parser.add_argument(\"--out\", default=None, help=\"Output filename.\")\n    parser.add_argument(\n        \"--config\",\n        default=\"{}\",\n        type=json.loads,\n        help=\"Algorithm-specific configuration (e.g. env, hyperparams). \"\n        \"Surpresses loading of configuration from checkpoint.\")\n    return parser\n\n\ndef run(args, parser):\n    config = args.config\n    if not config:\n        # Load configuration from file\n        config_dir = os.path.dirname(args.checkpoint)\n        config_path = os.path.join(config_dir, \"params.pkl\")\n        if not os.path.exists(config_path):\n            config_path = os.path.join(config_dir, \"../params.pkl\")\n        if not os.path.exists(config_path):\n            raise ValueError(\n                \"Could not find params.pkl in either the checkpoint dir or \"\n                \"its parent directory.\")\n        with open(config_path, 'rb') as f:\n            config = pickle.load(f)\n        if \"num_workers\" in config:\n            config[\"num_workers\"] = min(2, config[\"num_workers\"])\n\n    if not args.env:\n        if not config.get(\"env\"):\n            parser.error(\"the following arguments are required: --env\")\n        args.env = config.get(\"env\")\n\n    ray.init()\n\n    cls = get_agent_class(args.run)\n    agent = cls(env=args.env, config=config)\n    agent.restore(args.checkpoint)\n    num_steps = int(args.steps)\n    rollout(agent, args.env, num_steps, args.out, args.no_render)\n\n\ndef rollout(agent, env_name, num_steps, out=None, no_render=True):\n    if hasattr(agent, \"local_evaluator\"):\n        env = agent.local_evaluator.env\n        multiagent = agent.local_evaluator.multiagent\n        if multiagent:\n            policy_agent_mapping = agent.config[\"multiagent\"][\n                \"policy_mapping_fn\"]\n            mapping_cache = {}\n        policy_map = agent.local_evaluator.policy_map\n        state_init = {p: m.get_initial_state() for p, m in policy_map.items()}\n        use_lstm = {p: len(s) > 0 for p, s in state_init.items()}\n    else:\n        env = gym.make(env_name)\n        multiagent = False\n        use_lstm = {'default': False}\n\n    if out is not None:\n        rollouts = []\n    steps = 0\n    while steps < (num_steps or steps + 1):\n        if out is not None:\n            rollout = []\n        state = env.reset()\n        done = False\n        reward_total = 0.0\n        while not done and steps < (num_steps or steps + 1):\n            if multiagent:\n                action_dict = {}\n                for agent_id in state.keys():\n                    a_state = state[agent_id]\n                    if a_state is not None:\n                        policy_id = mapping_cache.setdefault(\n                            agent_id, policy_agent_mapping(agent_id))\n                        p_use_lstm = use_lstm[policy_id]\n                        if p_use_lstm:\n                            a_action, p_state_init, _ = agent.compute_action(\n                                a_state,\n                                state=state_init[policy_id],\n                                policy_id=policy_id)\n                            state_init[policy_id] = p_state_init\n                        else:\n                            a_action = agent.compute_action(\n                                a_state, policy_id=policy_id)\n                        action_dict[agent_id] = a_action\n                action = action_dict\n            else:\n                if use_lstm[\"default\"]:\n                    action, state_init, _ = agent.compute_action(\n                        state, state=state_init)\n                else:\n                    action = agent.compute_action(state)\n\n            next_state, reward, done, _ = env.step(action)\n\n            if multiagent:\n                done = done[\"__all__\"]\n                reward_total += sum(reward.values())\n            else:\n                reward_total += reward\n            if not no_render:\n                env.render()\n            if out is not None:\n                rollout.append([state, action, next_state, reward, done])\n            steps += 1\n            state = next_state\n        if out is not None:\n            rollouts.append(rollout)\n        print(\"Episode reward\", reward_total)\n\n    if out is not None:\n        pickle.dump(rollouts, open(out, \"wb\"))\n\n\nif __name__ == \"__main__\":\n    parser = create_parser()\n    args = parser.parse_args()\n    run(args, parser)\n", "path": "python/ray/rllib/rollout.py"}]}