We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
[rllib] MARWIL tuned cartpole example (and my own experiments) produce nan rewards only.
<!--Please include [tune], [rllib], [autoscaler] etc. in the issue title if relevant-->

### What is the problem? + Reproduction

I have a custom example that produces offline data and picks it up with MARWIL for training. I observed that I get `nan` reward values for my example every time, so I went a step back and used your cartpole example:

https://github.com/ray-project/ray/blob/cd5a207d69cdaf05b47d956c18e89d928585eec7/rllib/tuned_examples/marwil/cartpole-marwil.yaml

I'm following the exact steps there, i.e. first run 

```
./train.py --run=PPO --env=CartPole-v0 \
    --stop='{"timesteps_total": 50000}' \
    --config='{"output": "/tmp/out", "batch_mode": "complete_episodes"}'
```

followed by 

```
rllib train -f cartpole-marwil.yaml
```

I did this both on my currently preferred stable version `0.8.5`, as well as on the `0.9.0.dev0` wheel. The result is this:

```
== Status ==
Memory usage on this node: 19.4/32.0 GiB
Using FIFO scheduling algorithm.
Resources requested: 0/12 CPUs, 0/0 GPUs, 0.0/9.96 GiB heap, 0.0/3.42 GiB objects
Result logdir: /Users/maxpumperla/ray_results/cartpole-marwil
Number of trials: 2 (2 TERMINATED)
+--------------------------------+------------+-------+--------+--------+------------------+--------+----------+
| Trial name                     | status     | loc   |   beta |   iter |   total time (s) |     ts |   reward |
|--------------------------------+------------+-------+--------+--------+------------------+--------+----------|
| MARWIL_CartPole-v0_7af06_00000 | TERMINATED |       |      0 |   2206 |          58.5661 | 500007 |      nan |
| MARWIL_CartPole-v0_7af06_00001 | TERMINATED |       |      1 |   2248 |          58.6117 | 500286 |      nan |
+--------------------------------+------------+-------+--------+--------+------------------+--------+----------+
```

Also, I've noticed that your MARWIL unit test is a pure smoke test and doesn't check reward values, but I didn't run that locally. Maybe it produces nan values as well.

In any case I'd appreciate any input here, as we'd love to use MARWIL for our "real" use case, in which we see the same behaviour.

--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `rllib/examples/custom_loss.py`
Content:
```
1 """Example of using custom_loss() with an imitation learning loss.
2 
3 The default input file is too small to learn a good policy, but you can
4 generate new experiences for IL training as follows:
5 
6 To generate experiences:
7 $ ./train.py --run=PG --config='{"output": "/tmp/cartpole"}' --env=CartPole-v0
8 
9 To train on experiences with joint PG + IL loss:
10 $ python custom_loss.py --input-files=/tmp/cartpole
11 """
12 
13 import argparse
14 from pathlib import Path
15 import os
16 
17 import ray
18 from ray import tune
19 from ray.rllib.examples.models.custom_loss_model import CustomLossModel, \
20     TorchCustomLossModel
21 from ray.rllib.models import ModelCatalog
22 from ray.rllib.utils.framework import try_import_tf
23 
24 tf1, tf, tfv = try_import_tf()
25 
26 parser = argparse.ArgumentParser()
27 parser.add_argument("--torch", action="store_true")
28 parser.add_argument("--stop-iters", type=int, default=200)
29 parser.add_argument(
30     "--input-files",
31     type=str,
32     default=os.path.join(
33         os.path.dirname(os.path.abspath(__file__)),
34         "../tests/data/cartpole_small"))
35 
36 if __name__ == "__main__":
37     ray.init()
38     args = parser.parse_args()
39 
40     # Bazel makes it hard to find files specified in `args` (and `data`).
41     # Look for them here.
42     if not os.path.exists(args.input_files):
43         # This script runs in the ray/rllib/examples dir.
44         rllib_dir = Path(__file__).parent.parent
45         input_dir = rllib_dir.absolute().joinpath(args.input_files)
46         args.input_files = str(input_dir)
47 
48     ModelCatalog.register_custom_model(
49         "custom_loss", TorchCustomLossModel if args.torch else CustomLossModel)
50 
51     config = {
52         "env": "CartPole-v0",
53         "num_workers": 0,
54         "model": {
55             "custom_model": "custom_loss",
56             "custom_model_config": {
57                 "input_files": args.input_files,
58             },
59         },
60         "framework": "torch" if args.torch else "tf",
61     }
62 
63     stop = {
64         "training_iteration": args.stop_iters,
65     }
66 
67     tune.run("PG", config=config, stop=stop)
68
```
--- 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/rllib/examples/custom_loss.py b/rllib/examples/custom_loss.py
--- a/rllib/examples/custom_loss.py
+++ b/rllib/examples/custom_loss.py
@@ -31,7 +31,7 @@
     type=str,
     default=os.path.join(
         os.path.dirname(os.path.abspath(__file__)),
-        "../tests/data/cartpole_small"))
+        "../tests/data/cartpole/small"))
 
 if __name__ == "__main__":
     ray.init()

{"golden_diff": "diff --git a/rllib/examples/custom_loss.py b/rllib/examples/custom_loss.py\n--- a/rllib/examples/custom_loss.py\n+++ b/rllib/examples/custom_loss.py\n@@ -31,7 +31,7 @@\n     type=str,\n     default=os.path.join(\n         os.path.dirname(os.path.abspath(__file__)),\n-        \"../tests/data/cartpole_small\"))\n+        \"../tests/data/cartpole/small\"))\n \n if __name__ == \"__main__\":\n     ray.init()\n", "issue": "[rllib] MARWIL tuned cartpole example (and my own experiments) produce nan rewards only.\n<!--Please include [tune], [rllib], [autoscaler] etc. in the issue title if relevant-->\r\n\r\n### What is the problem? + Reproduction\r\n\r\nI have a custom example that produces offline data and picks it up with MARWIL for training. I observed that I get `nan` reward values for my example every time, so I went a step back and used your cartpole example:\r\n\r\nhttps://github.com/ray-project/ray/blob/cd5a207d69cdaf05b47d956c18e89d928585eec7/rllib/tuned_examples/marwil/cartpole-marwil.yaml\r\n\r\nI'm following the exact steps there, i.e. first run \r\n\r\n```\r\n./train.py --run=PPO --env=CartPole-v0 \\\r\n    --stop='{\"timesteps_total\": 50000}' \\\r\n    --config='{\"output\": \"/tmp/out\", \"batch_mode\": \"complete_episodes\"}'\r\n```\r\n\r\nfollowed by \r\n\r\n```\r\nrllib train -f cartpole-marwil.yaml\r\n```\r\n\r\nI did this both on my currently preferred stable version `0.8.5`, as well as on the `0.9.0.dev0` wheel. The result is this:\r\n\r\n```\r\n== Status ==\r\nMemory usage on this node: 19.4/32.0 GiB\r\nUsing FIFO scheduling algorithm.\r\nResources requested: 0/12 CPUs, 0/0 GPUs, 0.0/9.96 GiB heap, 0.0/3.42 GiB objects\r\nResult logdir: /Users/maxpumperla/ray_results/cartpole-marwil\r\nNumber of trials: 2 (2 TERMINATED)\r\n+--------------------------------+------------+-------+--------+--------+------------------+--------+----------+\r\n| Trial name                     | status     | loc   |   beta |   iter |   total time (s) |     ts |   reward |\r\n|--------------------------------+------------+-------+--------+--------+------------------+--------+----------|\r\n| MARWIL_CartPole-v0_7af06_00000 | TERMINATED |       |      0 |   2206 |          58.5661 | 500007 |      nan |\r\n| MARWIL_CartPole-v0_7af06_00001 | TERMINATED |       |      1 |   2248 |          58.6117 | 500286 |      nan |\r\n+--------------------------------+------------+-------+--------+--------+------------------+--------+----------+\r\n```\r\n\r\nAlso, I've noticed that your MARWIL unit test is a pure smoke test and doesn't check reward values, but I didn't run that locally. Maybe it produces nan values as well.\r\n\r\nIn any case I'd appreciate any input here, as we'd love to use MARWIL for our \"real\" use case, in which we see the same behaviour.\n", "before_files": [{"content": "\"\"\"Example of using custom_loss() with an imitation learning loss.\n\nThe default input file is too small to learn a good policy, but you can\ngenerate new experiences for IL training as follows:\n\nTo generate experiences:\n$ ./train.py --run=PG --config='{\"output\": \"/tmp/cartpole\"}' --env=CartPole-v0\n\nTo train on experiences with joint PG + IL loss:\n$ python custom_loss.py --input-files=/tmp/cartpole\n\"\"\"\n\nimport argparse\nfrom pathlib import Path\nimport os\n\nimport ray\nfrom ray import tune\nfrom ray.rllib.examples.models.custom_loss_model import CustomLossModel, \\\n    TorchCustomLossModel\nfrom ray.rllib.models import ModelCatalog\nfrom ray.rllib.utils.framework import try_import_tf\n\ntf1, tf, tfv = try_import_tf()\n\nparser = argparse.ArgumentParser()\nparser.add_argument(\"--torch\", action=\"store_true\")\nparser.add_argument(\"--stop-iters\", type=int, default=200)\nparser.add_argument(\n    \"--input-files\",\n    type=str,\n    default=os.path.join(\n        os.path.dirname(os.path.abspath(__file__)),\n        \"../tests/data/cartpole_small\"))\n\nif __name__ == \"__main__\":\n    ray.init()\n    args = parser.parse_args()\n\n    # Bazel makes it hard to find files specified in `args` (and `data`).\n    # Look for them here.\n    if not os.path.exists(args.input_files):\n        # This script runs in the ray/rllib/examples dir.\n        rllib_dir = Path(__file__).parent.parent\n        input_dir = rllib_dir.absolute().joinpath(args.input_files)\n        args.input_files = str(input_dir)\n\n    ModelCatalog.register_custom_model(\n        \"custom_loss\", TorchCustomLossModel if args.torch else CustomLossModel)\n\n    config = {\n        \"env\": \"CartPole-v0\",\n        \"num_workers\": 0,\n        \"model\": {\n            \"custom_model\": \"custom_loss\",\n            \"custom_model_config\": {\n                \"input_files\": args.input_files,\n            },\n        },\n        \"framework\": \"torch\" if args.torch else \"tf\",\n    }\n\n    stop = {\n        \"training_iteration\": args.stop_iters,\n    }\n\n    tune.run(\"PG\", config=config, stop=stop)\n", "path": "rllib/examples/custom_loss.py"}], "after_files": [{"content": "\"\"\"Example of using custom_loss() with an imitation learning loss.\n\nThe default input file is too small to learn a good policy, but you can\ngenerate new experiences for IL training as follows:\n\nTo generate experiences:\n$ ./train.py --run=PG --config='{\"output\": \"/tmp/cartpole\"}' --env=CartPole-v0\n\nTo train on experiences with joint PG + IL loss:\n$ python custom_loss.py --input-files=/tmp/cartpole\n\"\"\"\n\nimport argparse\nfrom pathlib import Path\nimport os\n\nimport ray\nfrom ray import tune\nfrom ray.rllib.examples.models.custom_loss_model import CustomLossModel, \\\n    TorchCustomLossModel\nfrom ray.rllib.models import ModelCatalog\nfrom ray.rllib.utils.framework import try_import_tf\n\ntf1, tf, tfv = try_import_tf()\n\nparser = argparse.ArgumentParser()\nparser.add_argument(\"--torch\", action=\"store_true\")\nparser.add_argument(\"--stop-iters\", type=int, default=200)\nparser.add_argument(\n    \"--input-files\",\n    type=str,\n    default=os.path.join(\n        os.path.dirname(os.path.abspath(__file__)),\n        \"../tests/data/cartpole/small\"))\n\nif __name__ == \"__main__\":\n    ray.init()\n    args = parser.parse_args()\n\n    # Bazel makes it hard to find files specified in `args` (and `data`).\n    # Look for them here.\n    if not os.path.exists(args.input_files):\n        # This script runs in the ray/rllib/examples dir.\n        rllib_dir = Path(__file__).parent.parent\n        input_dir = rllib_dir.absolute().joinpath(args.input_files)\n        args.input_files = str(input_dir)\n\n    ModelCatalog.register_custom_model(\n        \"custom_loss\", TorchCustomLossModel if args.torch else CustomLossModel)\n\n    config = {\n        \"env\": \"CartPole-v0\",\n        \"num_workers\": 0,\n        \"model\": {\n            \"custom_model\": \"custom_loss\",\n            \"custom_model_config\": {\n                \"input_files\": args.input_files,\n            },\n        },\n        \"framework\": \"torch\" if args.torch else \"tf\",\n    }\n\n    stop = {\n        \"training_iteration\": args.stop_iters,\n    }\n\n    tune.run(\"PG\", config=config, stop=stop)\n", "path": "rllib/examples/custom_loss.py"}]}