[RLlib] Examples folder restructuring (Model examples; final part). (#8278)

- This PR completes any previously missing PyTorch Model counterparts to TFModels in examples/models.
- It also makes sure, all example scripts in the rllib/examples folder are tested for both frameworks and learn the given task (this is often currently not checked) using a --as-test flag in connection with a --stop-reward.

rllib/examples/two_trainer_workflow.py

@@ -13,8 +13,10 @@ from ray import tune
 from ray.rllib.agents.trainer_template import build_trainer
 from ray.rllib.agents.dqn.dqn import DEFAULT_CONFIG as DQN_CONFIG
 from ray.rllib.agents.dqn.dqn_tf_policy import DQNTFPolicy
+from ray.rllib.agents.dqn.dqn_torch_policy import DQNTorchPolicy
 from ray.rllib.agents.ppo.ppo import DEFAULT_CONFIG as PPO_CONFIG
 from ray.rllib.agents.ppo.ppo_tf_policy import PPOTFPolicy
+from ray.rllib.agents.ppo.ppo_torch_policy import PPOTorchPolicy
 from ray.rllib.evaluation.worker_set import WorkerSet
 from ray.rllib.execution.common import _get_shared_metrics
 from ray.rllib.execution.concurrency_ops import Concurrently
@@ -25,10 +27,16 @@ from ray.rllib.execution.replay_ops import StoreToReplayBuffer, Replay
 from ray.rllib.execution.train_ops import TrainOneStep, UpdateTargetNetwork
 from ray.rllib.examples.env.multi_agent import MultiAgentCartPole
 from ray.rllib.optimizers.async_replay_optimizer import LocalReplayBuffer
+from ray.rllib.utils.test_utils import check_learning_achieved
 from ray.tune.registry import register_env
 
 parser = argparse.ArgumentParser()
-parser.add_argument("--num-iters", type=int, default=20)
+parser.add_argument("--as-test", action="store_true")
+parser.add_argument("--torch", action="store_true")
+parser.add_argument("--mixed-torch-tf", action="store_true")
+parser.add_argument("--stop-iters", type=int, default=20)
+parser.add_argument("--stop-reward", type=float, default=150.0)
+parser.add_argument("--stop-timesteps", type=int, default=100000)
 
 
 def custom_training_workflow(workers: WorkerSet, config: dict):
@@ -90,6 +98,9 @@ def custom_training_workflow(workers: WorkerSet, config: dict):
 
 if __name__ == "__main__":
     args = parser.parse_args()
+    assert not (args.torch and args.mixed_torch_tf),\
+        "Use either --torch or --mixed-torch-tf, not both!"
+
     ray.init()
 
     # Simple environment with 4 independent cartpole entities
@@ -102,8 +113,10 @@ if __name__ == "__main__":
     # Note that since the trainer below does not include a default policy or
     # policy configs, we have to explicitly set it in the multiagent config:
     policies = {
-        "ppo_policy": (PPOTFPolicy, obs_space, act_space, PPO_CONFIG),
-        "dqn_policy": (DQNTFPolicy, obs_space, act_space, DQN_CONFIG),
+        "ppo_policy": (PPOTorchPolicy if args.torch or args.mixed_torch_tf else
+                       PPOTFPolicy, obs_space, act_space, PPO_CONFIG),
+        "dqn_policy": (DQNTorchPolicy if args.torch else DQNTFPolicy,
+                       obs_space, act_space, DQN_CONFIG),
     }
 
     def policy_mapping_fn(agent_id):
@@ -117,16 +130,27 @@ if __name__ == "__main__":
         default_policy=None,
         execution_plan=custom_training_workflow)
 
-    tune.run(
-        MyTrainer,
-        stop={"training_iteration": args.num_iters},
-        config={
-            "rollout_fragment_length": 50,
-            "num_workers": 0,
-            "env": "multi_agent_cartpole",
-            "multiagent": {
-                "policies": policies,
-                "policy_mapping_fn": policy_mapping_fn,
-                "policies_to_train": ["dqn_policy", "ppo_policy"],
-            },
-        })
+    config = {
+        "rollout_fragment_length": 50,
+        "num_workers": 0,
+        "env": "multi_agent_cartpole",
+        "multiagent": {
+            "policies": policies,
+            "policy_mapping_fn": policy_mapping_fn,
+            "policies_to_train": ["dqn_policy", "ppo_policy"],
+        },
+        "use_pytorch": args.torch,
+    }
+
+    stop = {
+        "training_iteration": args.stop_iters,
+        "timesteps_total": args.stop_timesteps,
+        "episode_reward_mean": args.stop_reward,
+    }
+
+    results = tune.run(MyTrainer, config=config, stop=stop)
+
+    if args.as_test:
+        check_learning_achieved(results, args.stop_reward)
+
+    ray.shutdown()