Esempio n. 1
0
class A3CAgent(Agent):
    _agent_name = "A3C"
    _default_config = DEFAULT_CONFIG
    _allow_unknown_subkeys = ["model", "optimizer", "env_config"]

    @classmethod
    def default_resource_request(cls, config):
        cf = dict(cls._default_config, **config)
        return Resources(
            cpu=1,
            gpu=0,
            extra_cpu=cf["num_workers"],
            extra_gpu=cf["use_gpu_for_workers"] and cf["num_workers"] or 0)

    def _init(self):
        self.policy_cls = get_policy_cls(self.config)

        if self.config["use_pytorch"]:
            session_creator = None
        else:
            import tensorflow as tf

            def session_creator():
                return tf.Session(
                    config=tf.ConfigProto(
                        intra_op_parallelism_threads=1,
                        inter_op_parallelism_threads=1,
                        gpu_options=tf.GPUOptions(allow_growth=True)))

        remote_cls = CommonPolicyEvaluator.as_remote(
            num_gpus=1 if self.config["use_gpu_for_workers"] else 0)
        self.local_evaluator = CommonPolicyEvaluator(
            self.env_creator, self.policy_cls,
            batch_steps=self.config["batch_size"],
            batch_mode="truncate_episodes",
            tf_session_creator=session_creator,
            env_config=self.config["env_config"],
            model_config=self.config["model"], policy_config=self.config,
            num_envs=self.config["num_envs"])
        self.remote_evaluators = [
            remote_cls.remote(
                self.env_creator, self.policy_cls,
                batch_steps=self.config["batch_size"],
                batch_mode="truncate_episodes", sample_async=True,
                tf_session_creator=session_creator,
                env_config=self.config["env_config"],
                model_config=self.config["model"], policy_config=self.config,
                num_envs=self.config["num_envs"])
            for i in range(self.config["num_workers"])]

        self.optimizer = AsyncOptimizer(
            self.config["optimizer"], self.local_evaluator,
            self.remote_evaluators)

    def _train(self):
        self.optimizer.step()
        FilterManager.synchronize(
            self.local_evaluator.filters, self.remote_evaluators)
        return collect_metrics(self.local_evaluator, self.remote_evaluators)

    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 = {
            "remote_state": agent_state,
            "local_state": self.local_evaluator.save()
        }
        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"))
        ray.get([
            a.restore.remote(o)
            for a, o in zip(self.remote_evaluators, extra_data["remote_state"])
        ])
        self.local_evaluator.restore(extra_data["local_state"])

    def compute_action(self, observation, state=None):
        if state is None:
            state = []
        obs = self.local_evaluator.obs_filter(observation, update=False)
        return self.local_evaluator.for_policy(
            lambda p: p.compute_single_action(
                obs, state, is_training=False)[0])
Esempio n. 2
0
class PPOAgent(Agent):
    _agent_name = "PPO"
    _default_config = DEFAULT_CONFIG
    _default_policy_graph = PPOTFPolicyGraph

    @classmethod
    def default_resource_request(cls, config):
        cf = dict(cls._default_config, **config)
        return Resources(
            cpu=1,
            gpu=len([d for d in cf["devices"] if "gpu" in d.lower()]),
            extra_cpu=cf["num_cpus_per_worker"] * cf["num_workers"],
            extra_gpu=cf["num_gpus_per_worker"] * cf["num_workers"])

    def _init(self):
        def session_creator():
            return tf.Session(config=tf.ConfigProto(
                **self.config["tf_session_args"]))

        self.local_evaluator = CommonPolicyEvaluator(
            self.env_creator,
            self._default_policy_graph,
            tf_session_creator=session_creator,
            batch_mode="complete_episodes",
            observation_filter=self.config["observation_filter"],
            env_config=self.config["env_config"],
            model_config=self.config["model"],
            policy_config=self.config)
        RemoteEvaluator = CommonPolicyEvaluator.as_remote(
            num_cpus=self.config["num_cpus_per_worker"],
            num_gpus=self.config["num_gpus_per_worker"])
        self.remote_evaluators = [
            RemoteEvaluator.remote(
                self.env_creator,
                self._default_policy_graph,
                batch_mode="complete_episodes",
                observation_filter=self.config["observation_filter"],
                env_config=self.config["env_config"],
                model_config=self.config["model"],
                policy_config=self.config)
            for _ in range(self.config["num_workers"])
        ]

        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)

        # TODO(rliaw): Push into Policy Graph
        with self.local_evaluator.tf_sess.graph.as_default():
            self.saver = tf.train.Saver()

    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 = self.saver.save(self.local_evaluator.tf_sess,
                                          os.path.join(checkpoint_dir,
                                                       "checkpoint"),
                                          global_step=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):
        self.saver.restore(self.local_evaluator.tf_sess, 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])
        ])

    def compute_action(self, observation, state=None):
        if state is None:
            state = []
        obs = self.local_evaluator.filters["default"](observation,
                                                      update=False)
        return self.local_evaluator.for_policy(
            lambda p: p.compute_single_action(obs, state, is_training=False)[0
                                                                             ])
Esempio n. 3
0
class DQNAgent(Agent):
    _agent_name = "DQN"
    _default_config = DEFAULT_CONFIG
    _policy_graph = DQNPolicyGraph

    @classmethod
    def default_resource_request(cls, config):
        cf = dict(cls._default_config, **config)
        return Resources(
            cpu=1,
            gpu=cf["gpu"] and 1 or 0,
            extra_cpu=cf["num_cpus_per_worker"] * cf["num_workers"],
            extra_gpu=cf["num_gpus_per_worker"] * cf["num_workers"])

    def _init(self):
        adjusted_batch_size = (self.config["sample_batch_size"] +
                               self.config["n_step"] - 1)
        self.local_evaluator = CommonPolicyEvaluator(
            self.env_creator,
            self.config["multiagent"]["policy_graphs"] or self._policy_graph,
            policy_mapping_fn=self.config["multiagent"]["policy_mapping_fn"],
            batch_steps=adjusted_batch_size,
            batch_mode="truncate_episodes",
            preprocessor_pref="deepmind",
            compress_observations=True,
            env_config=self.config["env_config"],
            model_config=self.config["model"],
            policy_config=self.config,
            num_envs=self.config["num_envs"])
        remote_cls = CommonPolicyEvaluator.as_remote(
            num_cpus=self.config["num_cpus_per_worker"],
            num_gpus=self.config["num_gpus_per_worker"])
        self.remote_evaluators = [
            remote_cls.remote(self.env_creator,
                              self._policy_graph,
                              batch_steps=adjusted_batch_size,
                              batch_mode="truncate_episodes",
                              preprocessor_pref="deepmind",
                              compress_observations=True,
                              env_config=self.config["env_config"],
                              model_config=self.config["model"],
                              policy_config=self.config,
                              num_envs=self.config["num_envs"],
                              worker_index=i + 1)
            for i in range(self.config["num_workers"])
        ]

        self.exploration0 = self._make_exploration_schedule(0)
        self.explorations = [
            self._make_exploration_schedule(i)
            for i in range(self.config["num_workers"])
        ]

        for k in OPTIMIZER_SHARED_CONFIGS:
            if k not in self.config["optimizer_config"]:
                self.config["optimizer_config"][k] = self.config[k]

        self.optimizer = getattr(optimizers, self.config["optimizer_class"])(
            self.config["optimizer_config"], self.local_evaluator,
            self.remote_evaluators)

        self.last_target_update_ts = 0
        self.num_target_updates = 0

    def _make_exploration_schedule(self, worker_index):
        # Use either a different `eps` per worker, or a linear schedule.
        if self.config["per_worker_exploration"]:
            assert self.config["num_workers"] > 1, \
                "This requires multiple workers"
            return ConstantSchedule(0.4**(
                1 + worker_index / float(self.config["num_workers"] - 1) * 7))
        return LinearSchedule(
            schedule_timesteps=int(self.config["exploration_fraction"] *
                                   self.config["schedule_max_timesteps"]),
            initial_p=1.0,
            final_p=self.config["exploration_final_eps"])

    @property
    def global_timestep(self):
        return self.optimizer.num_steps_sampled

    def update_target_if_needed(self):
        if self.global_timestep - self.last_target_update_ts > \
                self.config["target_network_update_freq"]:
            self.local_evaluator.foreach_policy(lambda p, _: p.update_target())
            self.last_target_update_ts = self.global_timestep
            self.num_target_updates += 1

    def _train(self):
        start_timestep = self.global_timestep

        while (self.global_timestep - start_timestep <
               self.config["timesteps_per_iteration"]):
            self.optimizer.step()
            self.update_target_if_needed()

        exp_vals = [self.exploration0.value(self.global_timestep)]
        self.local_evaluator.foreach_policy(
            lambda p, _: p.set_epsilon(exp_vals[0]))
        for i, e in enumerate(self.remote_evaluators):
            exp_val = self.explorations[i].value(self.global_timestep)
            e.foreach_policy.remote(lambda p, _: p.set_epsilon(exp_val))
            exp_vals.append(exp_val)

        result = collect_metrics(self.local_evaluator, self.remote_evaluators)
        return result._replace(info=dict(
            {
                "min_exploration": min(exp_vals),
                "max_exploration": max(exp_vals),
                "num_target_updates": self.num_target_updates,
            }, **self.optimizer.stats()))

    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))
        extra_data = [
            self.local_evaluator.save(),
            ray.get([e.save.remote() for e in self.remote_evaluators]),
            self.optimizer.save(), self.num_target_updates,
            self.last_target_update_ts
        ]
        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([
            e.restore.remote(d)
            for (d, e) in zip(extra_data[1], self.remote_evaluators)
        ])
        self.optimizer.restore(extra_data[2])
        self.num_target_updates = extra_data[3]
        self.last_target_update_ts = extra_data[4]

    def compute_action(self, observation, state=None):
        if state is None:
            state = []
        return self.local_evaluator.for_policy(
            lambda p: p.compute_single_action(
                observation, state, is_training=False)[0])