[rllib] Document "v2" APIs (#2316)

* re * wip * wip * a3c working * torch support * pg works * lint * rm v2 * consumer id * clean up pg * clean up more * fix python 2.7 * tf session management * docs * dqn wip * fix compile * dqn * apex runs * up * impotrs * ddpg * quotes * fix tests * fix last r * fix tests * lint * pass checkpoint restore * kwar * nits * policy graph * fix yapf * com * class * pyt * vectorization * update * test cpe * unit test * fix ddpg2 * changes * wip * args * faster test * common * fix * add alg option * batch mode and policy serving * multi serving test * todo * wip * serving test * doc async env * num envs * comments * thread * remove init hook * update * fix ppo * comments1 * fix * updates * add jenkins tests * fix * fix pytorch * fix * fixes * fix a3c policy * fix squeeze * fix trunc on apex * fix squeezing for real * update * remove horizon test for now * multiagent wip * update * fix race condition * fix ma * t * doc * st * wip * example * wip * working * cartpole * wip * batch wip * fix bug * make other_batches None default * working * debug * nit * warn * comments * fix ppo * fix obs filter * update * wip * tf * update * fix * cleanup * cleanup * spacing * model * fix * dqn * fix ddpg * doc * keep names * update * fix * com * docs * clarify model outputs * Update torch_policy_graph.py * fix obs filter * pass thru worker index * fix * rename * vlad torch comments * fix log action * debug name * fix lstm * remove unused ddpg net * remove conv net * revert lstm * wip * wip * cast * wip * works * fix a3c * works * lstm util test * doc * clean up * update * fix lstm check * move to end * fix sphinx * fix cmd * remove bad doc * envs * vec * doc prep * models * rl * alg * up * clarify * copy * async sa * fix * comments * fix a3c conf * tune lstm * fix reshape * fix * back to 16 * tuned a3c update * update * tuned * optional * merge * wip * fix up * move pg class * rename env * wip * update * tip * alg * readme * fix catalog * readme * doc * context * remove prep * comma * add env * link to paper * paper * update * rnn * update * wip * clean up ev creation * fix * fix * fix * fix lint * up * no comma * ma * Update run_multi_node_tests.sh * fix * sphinx is stupid * sphinx is stupid * clarify torch graph * no horizon * fix config * sb * Update test_optimizers.py
2026-09-10 12:38:43 +08:00 · 2018-07-01 00:05:08 -07:00
parent 762bdf646e
commit 8aa56c12e6
113 changed files with 1148 additions and 1141 deletions
@@ -0,0 +1,140 @@
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import os
+import numpy as np
+import pickle
+
+import ray
+from ray.rllib.agents import Agent, with_common_config
+from ray.rllib.agents.ppo.ppo_tf_policy import PPOTFPolicyGraph
+from ray.rllib.evaluation.metrics import collect_metrics
+from ray.rllib.utils import FilterManager
+from ray.rllib.optimizers.multi_gpu_optimizer import LocalMultiGPUOptimizer
+from ray.tune.trial import Resources
+
+DEFAULT_CONFIG = with_common_config({
+    # If true, use the Generalized Advantage Estimator (GAE)
+    # with a value function, see https://arxiv.org/pdf/1506.02438.pdf.
+    "use_gae": True,
+    # GAE(lambda) parameter
+    "lambda": 1.0,
+    # Initial coefficient for KL divergence
+    "kl_coeff": 0.2,
+    # Number of timesteps collected for each SGD round
+    "timesteps_per_batch": 4000,
+    # Number of SGD iterations in each outer loop
+    "num_sgd_iter": 30,
+    # Stepsize of SGD
+    "sgd_stepsize": 5e-5,
+    # Total SGD batch size across all devices for SGD
+    "sgd_batchsize": 128,
+    # Coefficient of the value function loss
+    "vf_loss_coeff": 1.0,
+    # Coefficient of the entropy regularizer
+    "entropy_coeff": 0.0,
+    # PPO clip parameter
+    "clip_param": 0.3,
+    # Target value for KL divergence
+    "kl_target": 0.01,
+    # Number of GPUs to use for SGD
+    "num_gpus": 0,
+    # Whether to allocate GPUs for workers (if > 0).
+    "num_gpus_per_worker": 0,
+    # Whether to allocate CPUs for workers (if > 0).
+    "num_cpus_per_worker": 1,
+    # Whether to rollout "complete_episodes" or "truncate_episodes"
+    "batch_mode": "complete_episodes",
+    # Which observation filter to apply to the observation
+    "observation_filter": "MeanStdFilter",
+})
+
+
+class PPOAgent(Agent):
+    """Multi-GPU optimized implementation of PPO in TensorFlow."""
+
+    _agent_name = "PPO"
+    _default_config = DEFAULT_CONFIG
+
+    @classmethod
+    def default_resource_request(cls, config):
+        cf = dict(cls._default_config, **config)
+        return Resources(
+            cpu=1,
+            gpu=cf["num_gpus"],
+            extra_cpu=cf["num_cpus_per_worker"] * cf["num_workers"],
+            extra_gpu=cf["num_gpus_per_worker"] * cf["num_workers"])
+
+    def _init(self):
+        self.local_evaluator = self.make_local_evaluator(
+            self.env_creator, PPOTFPolicyGraph)
+        self.remote_evaluators = self.make_remote_evaluators(
+            self.env_creator, PPOTFPolicyGraph, self.config["num_workers"],
+            {"num_cpus": self.config["num_cpus_per_worker"],
+             "num_gpus": self.config["num_gpus_per_worker"]})
+        self.optimizer = LocalMultiGPUOptimizer(
+            {"sgd_batch_size": self.config["sgd_batchsize"],
+             "sgd_stepsize": self.config["sgd_stepsize"],
+             "num_sgd_iter": self.config["num_sgd_iter"],
+             "timesteps_per_batch": self.config["timesteps_per_batch"]},
+            self.local_evaluator, self.remote_evaluators)
+
+    def _train(self):
+        def postprocess_samples(batch):
+            # Divide by the maximum of value.std() and 1e-4
+            # to guard against the case where all values are equal
+            value = batch["advantages"]
+            standardized = (value - value.mean()) / max(1e-4, value.std())
+            batch.data["advantages"] = standardized
+            batch.shuffle()
+            dummy = np.zeros_like(batch["advantages"])
+            if not self.config["use_gae"]:
+                batch.data["value_targets"] = dummy
+                batch.data["vf_preds"] = dummy
+        extra_fetches = self.optimizer.step(postprocess_fn=postprocess_samples)
+        kl = np.array(extra_fetches["kl"]).mean(axis=1)[-1]
+        total_loss = np.array(extra_fetches["total_loss"]).mean(axis=1)[-1]
+        policy_loss = np.array(extra_fetches["policy_loss"]).mean(axis=1)[-1]
+        vf_loss = np.array(extra_fetches["vf_loss"]).mean(axis=1)[-1]
+        entropy = np.array(extra_fetches["entropy"]).mean(axis=1)[-1]
+
+        newkl = self.local_evaluator.for_policy(lambda pi: pi.update_kl(kl))
+
+        info = {
+            "kl_divergence": kl,
+            "kl_coefficient": newkl,
+            "total_loss": total_loss,
+            "policy_loss": policy_loss,
+            "vf_loss": vf_loss,
+            "entropy": entropy,
+        }
+
+        FilterManager.synchronize(
+            self.local_evaluator.filters, self.remote_evaluators)
+        res = collect_metrics(self.local_evaluator, self.remote_evaluators)
+        res = res._replace(info=info)
+        return res
+
+    def _stop(self):
+        # workaround for https://github.com/ray-project/ray/issues/1516
+        for ev in self.remote_evaluators:
+            ev.__ray_terminate__.remote()
+
+    def _save(self, checkpoint_dir):
+        checkpoint_path = os.path.join(checkpoint_dir,
+                                       "checkpoint-{}".format(self.iteration))
+        agent_state = ray.get(
+            [a.save.remote() for a in self.remote_evaluators])
+        extra_data = [
+            self.local_evaluator.save(),
+            agent_state]
+        pickle.dump(extra_data, open(checkpoint_path + ".extra_data", "wb"))
+        return checkpoint_path
+
+    def _restore(self, checkpoint_path):
+        extra_data = pickle.load(open(checkpoint_path + ".extra_data", "rb"))
+        self.local_evaluator.restore(extra_data[0])
+        ray.get([
+            a.restore.remote(o)
+                for (a, o) in zip(self.remote_evaluators, extra_data[1])])