From 81314143ebce027d36f51e32bb281721b60a7100 Mon Sep 17 00:00:00 2001 From: Sven Mika Date: Tue, 7 Apr 2020 21:40:34 +0200 Subject: [PATCH] [RLlib] Use framework_iterator (add torch/eager/tf) to PPO and PG tests. (#7915) --- rllib/agents/pg/tests/test_pg.py | 115 +++++++++++++++---------------- 1 file changed, 57 insertions(+), 58 deletions(-) diff --git a/rllib/agents/pg/tests/test_pg.py b/rllib/agents/pg/tests/test_pg.py index e0558132c..b4b1c800f 100644 --- a/rllib/agents/pg/tests/test_pg.py +++ b/rllib/agents/pg/tests/test_pg.py @@ -3,12 +3,11 @@ import unittest import ray import ray.rllib.agents.pg as pg -from ray.rllib.agents.pg import PGTrainer from ray.rllib.evaluation.postprocessing import Postprocessing from ray.rllib.models.tf.tf_action_dist import Categorical from ray.rllib.models.torch.torch_action_dist import TorchCategorical from ray.rllib.policy.sample_batch import SampleBatch -from ray.rllib.utils import check, fc +from ray.rllib.utils import check, fc, framework_iterator class TestPG(unittest.TestCase): @@ -19,7 +18,7 @@ class TestPG(unittest.TestCase): ray.shutdown() def test_pg_exec_impl(ray_start_regular): - trainer = PGTrainer( + trainer = pg.PGTrainer( env="CartPole-v0", config={ "min_iter_time_s": 0, @@ -31,25 +30,17 @@ class TestPG(unittest.TestCase): """Test whether a PGTrainer can be built with both frameworks.""" config = pg.DEFAULT_CONFIG.copy() config["num_workers"] = 0 # Run locally. - - # tf. - trainer = pg.PGTrainer(config=config, env="CartPole-v0") - num_iterations = 2 - for i in range(num_iterations): - trainer.train() - # Torch. - config["use_pytorch"] = True - trainer = pg.PGTrainer(config=config, env="CartPole-v0") - for i in range(num_iterations): - trainer.train() + for _ in framework_iterator(config): + trainer = pg.PGTrainer(config=config, env="CartPole-v0") + for i in range(num_iterations): + trainer.train() def test_pg_loss_functions(self): """Tests the PG loss function math.""" config = pg.DEFAULT_CONFIG.copy() config["num_workers"] = 0 # Run locally. - config["eager"] = True config["gamma"] = 0.99 config["model"]["fcnet_hiddens"] = [10] config["model"]["fcnet_activation"] = "linear" @@ -60,57 +51,65 @@ class TestPG(unittest.TestCase): 0.4], [0.5, 0.6, 0.7, 0.8], [0.9, 1.0, 1.1, 1.2]]), SampleBatch.ACTIONS: np.array([0, 1, 1]), + SampleBatch.PREV_ACTIONS: np.array([1, 0, 1]), SampleBatch.REWARDS: np.array([1.0, 1.0, 1.0]), + SampleBatch.PREV_REWARDS: np.array([-1.0, -1.0, -1.0]), SampleBatch.DONES: np.array([False, False, True]) } - # tf. - trainer = pg.PGTrainer(config=config, env="CartPole-v0") - policy = trainer.get_policy() - vars = policy.model.trainable_variables() + for fw, sess in framework_iterator(config, session=True): + dist_cls = (Categorical if fw != "torch" else TorchCategorical) + trainer = pg.PGTrainer(config=config, env="CartPole-v0") + policy = trainer.get_policy() + vars = policy.model.trainable_variables() + if fw == "tf": + vars = policy.get_session().run(vars) - # Post-process (calculate simple (non-GAE) advantages) and attach to - # train_batch dict. - # A = [0.99^2 * 1.0 + 0.99 * 1.0 + 1.0, 0.99 * 1.0 + 1.0, 1.0] = - # [2.9701, 1.99, 1.0] - train_batch = pg.post_process_advantages(policy, train_batch) - # Check Advantage values. - check(train_batch[Postprocessing.ADVANTAGES], [2.9701, 1.99, 1.0]) + # Post-process (calculate simple (non-GAE) advantages) and attach + # to train_batch dict. + # A = [0.99^2 * 1.0 + 0.99 * 1.0 + 1.0, 0.99 * 1.0 + 1.0, 1.0] = + # [2.9701, 1.99, 1.0] + train_batch = pg.post_process_advantages(policy, train_batch) + if fw == "torch": + train_batch = policy._lazy_tensor_dict(train_batch) - # Actual loss results. - results = pg.pg_tf_loss( - policy, - policy.model, - dist_class=Categorical, - train_batch=train_batch) + # Check Advantage values. + check(train_batch[Postprocessing.ADVANTAGES], [2.9701, 1.99, 1.0]) - # Calculate expected results. - expected_logits = fc( - fc(train_batch[SampleBatch.CUR_OBS], vars[0].numpy(), - vars[1].numpy()), vars[2].numpy(), vars[3].numpy()) - expected_logp = Categorical(expected_logits, policy.model).logp( - train_batch[SampleBatch.ACTIONS]) - expected_loss = -np.mean( - expected_logp * train_batch[Postprocessing.ADVANTAGES]) - check(results.numpy(), expected_loss, decimals=4) + # Actual loss results. + if fw == "tf": + results = policy.get_session().run( + policy._loss, + feed_dict=policy._get_loss_inputs_dict( + train_batch, shuffle=False)) + else: + results = (pg.pg_tf_loss + if fw == "eager" else pg.pg_torch_loss)( + policy, + policy.model, + dist_class=dist_cls, + train_batch=train_batch) - # Torch. - config["use_pytorch"] = True - trainer = pg.PGTrainer(config=config, env="CartPole-v0") - policy = trainer.get_policy() - train_batch = policy._lazy_tensor_dict(train_batch) - results = pg.pg_torch_loss( - policy, - policy.model, - dist_class=TorchCategorical, - train_batch=train_batch) - expected_logits = policy.model.last_output() - expected_logp = TorchCategorical(expected_logits, policy.model).logp( - train_batch[SampleBatch.ACTIONS]) - expected_loss = -np.mean( - expected_logp.detach().numpy() * - train_batch[Postprocessing.ADVANTAGES].numpy()) - check(results.detach().numpy(), expected_loss, decimals=4) + # Calculate expected results. + expected_logits = fc( + fc(train_batch[SampleBatch.CUR_OBS], + vars[0], + vars[1], + framework=fw), + vars[2], + vars[3], + framework=fw) + expected_logp = dist_cls(expected_logits, policy.model).logp( + train_batch[SampleBatch.ACTIONS]) + if sess: + expected_logp = sess.run(expected_logp) + else: + expected_logp = expected_logp.numpy() + expected_loss = -np.mean( + expected_logp * + (train_batch[Postprocessing.ADVANTAGES] if fw != "torch" else + train_batch[Postprocessing.ADVANTAGES].numpy())) + check(results, expected_loss, decimals=4) if __name__ == "__main__":