We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
rllib rollout does not load the model automatically from params.json
<!--
General questions should be asked on the mailing list ray-dev@googlegroups.com.
Questions about how to use Ray should be asked on
[StackOverflow](https://stackoverflow.com/questions/tagged/ray).

Before submitting an issue, please fill out the following form.
-->

### System information
- **OS Platform and Distribution (e.g., Linux Ubuntu 16.04)**: Linux 4.4.0-135-generic x86_64 
- **Python version**: Python 3.6.5 

<!--
You can obtain the Ray version with

python -c "import ray; print(ray.__version__)"
-->

### Describe the problem
<!-- Describe the problem clearly here. -->

### Source code / logs
<!-- Include any logs or source code that would be helpful to diagnose the problem. If including tracebacks, please include the full traceback. Large logs and files should be attached. Try to provide a reproducible test case that is the bare minimum necessary to generate the problem. -->
rllib rollout does not load the model automatically from params.json for a simple 256x256x256x256 model. 
When I run rllib rollout without specifying --config with "model": {"fcnet_hiddens": [256, 256, 256, 256]} it fails with the following error:
```
    assert len(vector) == i, "Passed weight does not have the correct shape."
AssertionError: Passed weight does not have the correct shape.
```

--- 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.pkl")
77         if not os.path.exists(config_path):
78             config_path = os.path.join(config_dir, "../params.pkl")
79         if not os.path.exists(config_path):
80             raise ValueError(
81                 "Could not find params.pkl in either the checkpoint dir or "
82                 "its parent directory.")
83         with open(config_path, 'rb') as f:
84             config = pickle.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         multiagent = agent.local_evaluator.multiagent
106         if multiagent:
107             policy_agent_mapping = agent.config["multiagent"][
108                 "policy_mapping_fn"]
109             mapping_cache = {}
110         policy_map = agent.local_evaluator.policy_map
111         state_init = {p: m.get_initial_state() for p, m in policy_map.items()}
112         use_lstm = {p: len(s) > 0 for p, s in state_init.items()}
113     else:
114         env = gym.make(env_name)
115         multiagent = False
116         use_lstm = {'default': 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 multiagent:
129                 action_dict = {}
130                 for agent_id in state.keys():
131                     a_state = state[agent_id]
132                     if a_state is not None:
133                         policy_id = mapping_cache.setdefault(
134                             agent_id, policy_agent_mapping(agent_id))
135                         p_use_lstm = use_lstm[policy_id]
136                         if p_use_lstm:
137                             a_action, p_state_init, _ = agent.compute_action(
138                                 a_state,
139                                 state=state_init[policy_id],
140                                 policy_id=policy_id)
141                             state_init[policy_id] = p_state_init
142                         else:
143                             a_action = agent.compute_action(
144                                 a_state, policy_id=policy_id)
145                         action_dict[agent_id] = a_action
146                 action = action_dict
147             else:
148                 if use_lstm["default"]:
149                     action, state_init, _ = agent.compute_action(
150                         state, state=state_init)
151                 else:
152                     action = agent.compute_action(state)
153 
154             next_state, reward, done, _ = env.step(action)
155 
156             if multiagent:
157                 done = done["__all__"]
158                 reward_total += sum(reward.values())
159             else:
160                 reward_total += reward
161             if not no_render:
162                 env.render()
163             if out is not None:
164                 rollout.append([state, action, next_state, reward, done])
165             steps += 1
166             state = next_state
167         if out is not None:
168             rollouts.append(rollout)
169         print("Episode reward", reward_total)
170 
171     if out is not None:
172         pickle.dump(rollouts, open(out, "wb"))
173 
174 
175 if __name__ == "__main__":
176     parser = create_parser()
177     args = parser.parse_args()
178     run(args, parser)
179
```
--- 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()

```

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
@@ -12,6 +12,7 @@
 import gym
 import ray
 from ray.rllib.agents.registry import get_agent_class
+from ray.tune.util import merge_dicts
 
 EXAMPLE_USAGE = """
 Example Usage via RLlib CLI:
@@ -69,22 +70,23 @@
 
 
 def run(args, parser):
-    config = args.config
-    if not config:
-        # Load configuration from file
-        config_dir = os.path.dirname(args.checkpoint)
-        config_path = os.path.join(config_dir, "params.pkl")
-        if not os.path.exists(config_path):
-            config_path = os.path.join(config_dir, "../params.pkl")
-        if not os.path.exists(config_path):
+    config = {}
+    # Load configuration from file
+    config_dir = os.path.dirname(args.checkpoint)
+    config_path = os.path.join(config_dir, "params.pkl")
+    if not os.path.exists(config_path):
+        config_path = os.path.join(config_dir, "../params.pkl")
+    if not os.path.exists(config_path):
+        if not args.config:
             raise ValueError(
                 "Could not find params.pkl in either the checkpoint dir or "
                 "its parent directory.")
+    else:
         with open(config_path, 'rb') as f:
             config = pickle.load(f)
-        if "num_workers" in config:
-            config["num_workers"] = min(2, config["num_workers"])
-
+    if "num_workers" in config:
+        config["num_workers"] = min(2, config["num_workers"])
+    config = merge_dicts(config, args.config)
     if not args.env:
         if not config.get("env"):
             parser.error("the following arguments are required: --env")

{"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@@ -12,6 +12,7 @@\n import gym\n import ray\n from ray.rllib.agents.registry import get_agent_class\n+from ray.tune.util import merge_dicts\n \n EXAMPLE_USAGE = \"\"\"\n Example Usage via RLlib CLI:\n@@ -69,22 +70,23 @@\n \n \n def 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+    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+        if not args.config:\n             raise ValueError(\n                 \"Could not find params.pkl in either the checkpoint dir or \"\n                 \"its parent directory.\")\n+    else:\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 \"num_workers\" in config:\n+        config[\"num_workers\"] = min(2, config[\"num_workers\"])\n+    config = merge_dicts(config, args.config)\n     if not args.env:\n         if not config.get(\"env\"):\n             parser.error(\"the following arguments are required: --env\")\n", "issue": "rllib rollout does not load the model automatically from params.json\n<!--\r\nGeneral questions should be asked on the mailing list ray-dev@googlegroups.com.\r\nQuestions about how to use Ray should be asked on\r\n[StackOverflow](https://stackoverflow.com/questions/tagged/ray).\r\n\r\nBefore submitting an issue, please fill out the following form.\r\n-->\r\n\r\n### System information\r\n- **OS Platform and Distribution (e.g., Linux Ubuntu 16.04)**: Linux 4.4.0-135-generic x86_64 \r\n- **Python version**: Python 3.6.5 \r\n\r\n<!--\r\nYou can obtain the Ray version with\r\n\r\npython -c \"import ray; print(ray.__version__)\"\r\n-->\r\n\r\n### Describe the problem\r\n<!-- Describe the problem clearly here. -->\r\n\r\n### Source code / logs\r\n<!-- Include any logs or source code that would be helpful to diagnose the problem. If including tracebacks, please include the full traceback. Large logs and files should be attached. Try to provide a reproducible test case that is the bare minimum necessary to generate the problem. -->\r\nrllib rollout does not load the model automatically from params.json for a simple 256x256x256x256 model. \r\nWhen I run rllib rollout without specifying --config with \"model\": {\"fcnet_hiddens\": [256, 256, 256, 256]} it fails with the following error:\r\n```\r\n    assert len(vector) == i, \"Passed weight does not have the correct shape.\"\r\nAssertionError: Passed weight does not have the correct shape.\r\n```\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.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"}], "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\nfrom ray.tune.util import merge_dicts\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 = {}\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        if not args.config:\n            raise ValueError(\n                \"Could not find params.pkl in either the checkpoint dir or \"\n                \"its parent directory.\")\n    else:\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    config = merge_dicts(config, args.config)\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"}]}