[RLlib] Fix all example scripts to run on GPUs. (#11105)

rllib/examples/complex_struct_space.py

@@ -8,7 +8,9 @@ For PyTorch / TF eager mode, use the --torch and --eager flags.
 """
 
 import argparse
+import os
 
+import ray
 from ray import tune
 from ray.rllib.models import ModelCatalog
 from ray.rllib.examples.env.simple_rpg import SimpleRPG
@@ -17,9 +19,10 @@ from ray.rllib.examples.models.simple_rpg_model import CustomTorchRPGModel, \
 
 parser = argparse.ArgumentParser()
 parser.add_argument(
-    "--framework", choices=["tf", "tfe", "torch"], default="tf")
+    "--framework", choices=["tf2", "tf", "tfe", "torch"], default="tf2")
 
 if __name__ == "__main__":
+    ray.init()
     args = parser.parse_args()
     if args.framework == "torch":
         ModelCatalog.register_custom_model("my_model", CustomTorchRPGModel)
@@ -31,6 +34,8 @@ if __name__ == "__main__":
         "env": SimpleRPG,
         "rollout_fragment_length": 1,
         "train_batch_size": 2,
+        # Use GPUs iff `RLLIB_NUM_GPUS` env var set to > 0.
+        "num_gpus": int(os.environ.get("RLLIB_NUM_GPUS", "0")),
         "num_workers": 0,
         "model": {
             "custom_model": "my_model",