diff --git a/rllib/BUILD b/rllib/BUILD index 270005069..ab62877a3 100644 --- a/rllib/BUILD +++ b/rllib/BUILD @@ -348,7 +348,7 @@ py_test( "--env", "Pendulum-v0", "--run", "APEX_DDPG", "--stop", "'{\"training_iteration\": 1}'", - "--config", "'{\"num_workers\": 2, \"optimizer\": {\"num_replay_buffer_shards\": 1}, \"learning_starts\": 100, \"min_iter_time_s\": 1, \"batch_mode\": \"complete_episodes\", \"parameter_noise\": false}'", + "--config", "'{\"num_workers\": 2, \"optimizer\": {\"num_replay_buffer_shards\": 1}, \"learning_starts\": 100, \"min_iter_time_s\": 1, \"batch_mode\": \"complete_episodes\"}'", "--ray-num-cpus", "4", ] ) diff --git a/rllib/agents/ddpg/ddpg_policy.py b/rllib/agents/ddpg/ddpg_policy.py index b00ee4a0b..f0d54b4a6 100644 --- a/rllib/agents/ddpg/ddpg_policy.py +++ b/rllib/agents/ddpg/ddpg_policy.py @@ -74,6 +74,8 @@ class DDPGPostprocessing: class DDPGTFPolicy(DDPGPostprocessing, TFPolicy): def __init__(self, observation_space, action_space, config): + self.observation_space = observation_space + self.action_space = action_space config = dict(ray.rllib.agents.ddpg.ddpg.DEFAULT_CONFIG, **config) if not isinstance(action_space, Box): raise UnsupportedSpaceException( @@ -106,9 +108,11 @@ class DDPGTFPolicy(DDPGPostprocessing, TFPolicy): name="cur_obs") with tf.variable_scope(POLICY_SCOPE) as scope: - policy_out, self.policy_model = self._build_policy_network( - self.cur_observations, observation_space, action_space) + self._distribution_inputs, self.policy_model = \ + self._build_policy_network( + self.cur_observations, observation_space, action_space) self.policy_vars = scope_vars(scope.name) + self.model = self.policy_model # Noise vars for P network except for layer normalization vars if self.config["parameter_noise"]: @@ -117,15 +121,17 @@ class DDPGTFPolicy(DDPGPostprocessing, TFPolicy): ]) # Create exploration component. - self.exploration = self._create_exploration(action_space, config) + self.exploration = self._create_exploration() explore = tf.placeholder_with_default(True, (), name="is_exploring") - # Action outputs + # Action outputs. with tf.variable_scope(ACTION_SCOPE): self.output_actions, _ = self.exploration.get_exploration_action( - policy_out, Deterministic, self.policy_model, timestep, - explore) + action_distribution=Deterministic(self._distribution_inputs, + self.model), + timestep=timestep, + explore=explore) - # Replay inputs + # Replay inputs. self.obs_t = tf.placeholder( tf.float32, shape=(None, ) + observation_space.shape, @@ -289,6 +295,8 @@ class DDPGTFPolicy(DDPGPostprocessing, TFPolicy): loss_inputs=self.loss_inputs, update_ops=q_batchnorm_update_ops + policy_batchnorm_update_ops, explore=explore, + dist_inputs=self._distribution_inputs, + dist_class=Deterministic, timestep=timestep) self.sess.run(tf.global_variables_initializer()) diff --git a/rllib/agents/dqn/dqn_policy.py b/rllib/agents/dqn/dqn_policy.py index d18d540a3..a7c1dc46a 100644 --- a/rllib/agents/dqn/dqn_policy.py +++ b/rllib/agents/dqn/dqn_policy.py @@ -7,11 +7,11 @@ from ray.rllib.agents.dqn.distributional_q_model import DistributionalQModel from ray.rllib.agents.dqn.simple_q_policy import TargetNetworkMixin, \ ParameterNoiseMixin from ray.rllib.policy.sample_batch import SampleBatch +from ray.rllib.policy.tf_policy import LearningRateSchedule +from ray.rllib.policy.tf_policy_template import build_tf_policy from ray.rllib.models import ModelCatalog from ray.rllib.models.tf.tf_action_dist import Categorical from ray.rllib.utils.error import UnsupportedSpaceException -from ray.rllib.policy.tf_policy import LearningRateSchedule -from ray.rllib.policy.tf_policy_template import build_tf_policy from ray.rllib.utils.tf_ops import huber_loss, reduce_mean_ignore_inf, \ minimize_and_clip from ray.rllib.utils import try_import_tf @@ -202,88 +202,35 @@ def build_q_model(policy, obs_space, action_space, config): return policy.q_model -def get_log_likelihood(policy, q_model, actions, input_dict, obs_space, - action_space, config): - # Action Q network. - q_vals = _compute_q_values(policy, q_model, - input_dict[SampleBatch.CUR_OBS], obs_space, - action_space) +def get_distribution_inputs_and_class(policy, + q_model, + obs_batch, + *, + explore=True, + **kwargs): + q_vals = compute_q_values(policy, q_model, obs_batch, explore) q_vals = q_vals[0] if isinstance(q_vals, tuple) else q_vals - action_dist = Categorical(q_vals, q_model) - return action_dist.logp(actions) - -def sample_action_from_q_network(policy, q_model, input_dict, obs_space, - action_space, explore, config, timestep): - # Action Q network. - q_vals = _compute_q_values(policy, q_model, - input_dict[SampleBatch.CUR_OBS], obs_space, - action_space) - policy.q_values = q_vals[0] if isinstance(q_vals, tuple) else q_vals + policy.q_values = q_vals policy.q_func_vars = q_model.variables() - - policy.output_actions, policy.sampled_action_logp = \ - policy.exploration.get_exploration_action( - policy.q_values, Categorical, q_model, timestep, explore) - - # Noise vars for Q network except for layer normalization vars. - if config["parameter_noise"]: - _build_parameter_noise( - policy, - [var for var in policy.q_func_vars if "LayerNorm" not in var.name]) - policy.action_probs = tf.nn.softmax(policy.q_values) - - return policy.output_actions, policy.sampled_action_logp - - -def _build_parameter_noise(policy, pnet_params): - policy.parameter_noise_sigma_val = 1.0 - policy.parameter_noise_sigma = tf.get_variable( - initializer=tf.constant_initializer(policy.parameter_noise_sigma_val), - name="parameter_noise_sigma", - shape=(), - trainable=False, - dtype=tf.float32) - policy.parameter_noise = list() - # No need to add any noise on LayerNorm parameters - for var in pnet_params: - noise_var = tf.get_variable( - name=var.name.split(":")[0] + "_noise", - shape=var.shape, - initializer=tf.constant_initializer(.0), - trainable=False) - policy.parameter_noise.append(noise_var) - remove_noise_ops = list() - for var, var_noise in zip(pnet_params, policy.parameter_noise): - remove_noise_ops.append(tf.assign_add(var, -var_noise)) - policy.remove_noise_op = tf.group(*tuple(remove_noise_ops)) - generate_noise_ops = list() - for var_noise in policy.parameter_noise: - generate_noise_ops.append( - tf.assign( - var_noise, - tf.random_normal( - shape=var_noise.shape, - stddev=policy.parameter_noise_sigma))) - with tf.control_dependencies(generate_noise_ops): - add_noise_ops = list() - for var, var_noise in zip(pnet_params, policy.parameter_noise): - add_noise_ops.append(tf.assign_add(var, var_noise)) - policy.add_noise_op = tf.group(*tuple(add_noise_ops)) - policy.pi_distance = None + return policy.q_values, Categorical, [] # state-out def build_q_losses(policy, model, _, train_batch): config = policy.config # q network evaluation - q_t, q_logits_t, q_dist_t = _compute_q_values( - policy, policy.q_model, train_batch[SampleBatch.CUR_OBS], - policy.observation_space, policy.action_space) + q_t, q_logits_t, q_dist_t = compute_q_values( + policy, + policy.q_model, + train_batch[SampleBatch.CUR_OBS], + explore=False) # target q network evalution - q_tp1, q_logits_tp1, q_dist_tp1 = _compute_q_values( - policy, policy.target_q_model, train_batch[SampleBatch.NEXT_OBS], - policy.observation_space, policy.action_space) + q_tp1, q_logits_tp1, q_dist_tp1 = compute_q_values( + policy, + policy.target_q_model, + train_batch[SampleBatch.NEXT_OBS], + explore=False) policy.target_q_func_vars = policy.target_q_model.variables() # q scores for actions which we know were selected in the given state. @@ -297,10 +244,10 @@ def build_q_losses(policy, model, _, train_batch): # compute estimate of best possible value starting from state at t + 1 if config["double_q"]: q_tp1_using_online_net, q_logits_tp1_using_online_net, \ - q_dist_tp1_using_online_net = _compute_q_values( + q_dist_tp1_using_online_net = compute_q_values( policy, policy.q_model, train_batch[SampleBatch.NEXT_OBS], - policy.observation_space, policy.action_space) + explore=False) q_tp1_best_using_online_net = tf.argmax(q_tp1_using_online_net, 1) q_tp1_best_one_hot_selection = tf.one_hot(q_tp1_best_using_online_net, policy.action_space.n) @@ -362,10 +309,11 @@ def setup_late_mixins(policy, obs_space, action_space, config): TargetNetworkMixin.__init__(policy, obs_space, action_space, config) -def _compute_q_values(policy, model, obs, obs_space, action_space): +def compute_q_values(policy, model, obs, explore): config = policy.config + model_out, state = model({ - "obs": obs, + SampleBatch.CUR_OBS: obs, "is_training": policy._get_is_training_placeholder(), }, [], None) @@ -456,8 +404,7 @@ DQNTFPolicy = build_tf_policy( name="DQNTFPolicy", get_default_config=lambda: ray.rllib.agents.dqn.dqn.DEFAULT_CONFIG, make_model=build_q_model, - action_sampler_fn=sample_action_from_q_network, - log_likelihood_fn=get_log_likelihood, + action_distribution_fn=get_distribution_inputs_and_class, loss_fn=build_q_losses, stats_fn=build_q_stats, postprocess_fn=postprocess_nstep_and_prio, diff --git a/rllib/agents/dqn/simple_q_policy.py b/rllib/agents/dqn/simple_q_policy.py index 5d455ba75..31c623026 100644 --- a/rllib/agents/dqn/simple_q_policy.py +++ b/rllib/agents/dqn/simple_q_policy.py @@ -88,44 +88,34 @@ def build_q_models(policy, obs_space, action_space, config): return policy.q_model -def get_log_likelihood(policy, q_model, actions, input_dict, obs_space, - action_space, config): - # Action Q network. - q_vals = _compute_q_values(policy, q_model, - input_dict[SampleBatch.CUR_OBS], obs_space, - action_space) +def get_distribution_inputs_and_class(policy, + q_model, + obs_batch, + *, + explore=True, + **kwargs): + q_vals = compute_q_values(policy, q_model, obs_batch, explore) q_vals = q_vals[0] if isinstance(q_vals, tuple) else q_vals - action_dist = Categorical(q_vals, q_model) - return action_dist.logp(actions) - -def simple_sample_action_from_q_network(policy, q_model, input_dict, obs_space, - action_space, explore, config, - timestep): - # Action Q network. - q_vals = _compute_q_values(policy, q_model, - input_dict[SampleBatch.CUR_OBS], obs_space, - action_space) - policy.q_values = q_vals[0] if isinstance(q_vals, tuple) else q_vals + policy.q_values = q_vals policy.q_func_vars = q_model.variables() - - policy.output_actions, policy.sampled_action_logp = \ - policy.exploration.get_exploration_action( - policy.q_values, Categorical, q_model, timestep, explore) - - return policy.output_actions, policy.sampled_action_logp + return policy.q_values, Categorical, [] # state-outs def build_q_losses(policy, model, dist_class, train_batch): # q network evaluation - q_t = _compute_q_values(policy, policy.q_model, - train_batch[SampleBatch.CUR_OBS], - policy.observation_space, policy.action_space) + q_t = compute_q_values( + policy, + policy.q_model, + train_batch[SampleBatch.CUR_OBS], + explore=False) # target q network evalution - q_tp1 = _compute_q_values(policy, policy.target_q_model, - train_batch[SampleBatch.NEXT_OBS], - policy.observation_space, policy.action_space) + q_tp1 = compute_q_values( + policy, + policy.target_q_model, + train_batch[SampleBatch.NEXT_OBS], + explore=False) policy.target_q_func_vars = policy.target_q_model.variables() # q scores for actions which we know were selected in the given state. @@ -155,12 +145,12 @@ def build_q_losses(policy, model, dist_class, train_batch): return loss -def _compute_q_values(policy, model, obs, obs_space, action_space): - input_dict = { - "obs": obs, +def compute_q_values(policy, model, obs, explore): + model_out, _ = model({ + SampleBatch.CUR_OBS: obs, "is_training": policy._get_is_training_placeholder(), - } - model_out, _ = model(input_dict, [], None) + }, [], None) + return model.get_q_values(model_out) @@ -176,8 +166,7 @@ SimpleQPolicy = build_tf_policy( name="SimpleQPolicy", get_default_config=lambda: ray.rllib.agents.dqn.dqn.DEFAULT_CONFIG, make_model=build_q_models, - action_sampler_fn=simple_sample_action_from_q_network, - log_likelihood_fn=get_log_likelihood, + action_distribution_fn=get_distribution_inputs_and_class, loss_fn=build_q_losses, extra_action_fetches_fn=lambda policy: {"q_values": policy.q_values}, extra_learn_fetches_fn=lambda policy: {"td_error": policy.td_error}, diff --git a/rllib/agents/impala/vtrace_policy.py b/rllib/agents/impala/vtrace_policy.py index 1d5e1960d..6480556fa 100644 --- a/rllib/agents/impala/vtrace_policy.py +++ b/rllib/agents/impala/vtrace_policy.py @@ -12,7 +12,7 @@ from ray.rllib.models.tf.tf_action_dist import Categorical from ray.rllib.policy.sample_batch import SampleBatch from ray.rllib.policy.tf_policy_template import build_tf_policy from ray.rllib.policy.tf_policy import LearningRateSchedule, \ - EntropyCoeffSchedule, ACTION_LOGP + EntropyCoeffSchedule from ray.rllib.utils.explained_variance import explained_variance from ray.rllib.utils import try_import_tf @@ -20,8 +20,6 @@ tf = try_import_tf() logger = logging.getLogger(__name__) -BEHAVIOUR_LOGITS = "behaviour_logits" - class VTraceLoss: def __init__(self, @@ -171,8 +169,8 @@ def build_vtrace_loss(policy, model, dist_class, train_batch): actions = train_batch[SampleBatch.ACTIONS] dones = train_batch[SampleBatch.DONES] rewards = train_batch[SampleBatch.REWARDS] - behaviour_action_logp = train_batch[ACTION_LOGP] - behaviour_logits = train_batch[BEHAVIOUR_LOGITS] + behaviour_action_logp = train_batch[SampleBatch.ACTION_LOGP] + behaviour_logits = train_batch[SampleBatch.ACTION_DIST_INPUTS] unpacked_behaviour_logits = tf.split( behaviour_logits, output_hidden_shape, axis=1) unpacked_outputs = tf.split(model_out, output_hidden_shape, axis=1) @@ -253,10 +251,6 @@ def postprocess_trajectory(policy, return sample_batch -def add_behaviour_logits(policy): - return {BEHAVIOUR_LOGITS: policy.model.last_output()} - - def validate_config(policy, obs_space, action_space, config): if config["vtrace"] and not config["in_evaluation"]: assert config["batch_mode"] == "truncate_episodes", \ @@ -295,7 +289,6 @@ VTraceTFPolicy = build_tf_policy( postprocess_fn=postprocess_trajectory, optimizer_fn=choose_optimizer, gradients_fn=clip_gradients, - extra_action_fetches_fn=add_behaviour_logits, before_init=validate_config, before_loss_init=setup_mixins, mixins=[LearningRateSchedule, EntropyCoeffSchedule], diff --git a/rllib/agents/ppo/appo_policy.py b/rllib/agents/ppo/appo_policy.py index 2b0d9d30e..0662bf0f0 100644 --- a/rllib/agents/ppo/appo_policy.py +++ b/rllib/agents/ppo/appo_policy.py @@ -8,7 +8,7 @@ import gym from ray.rllib.agents.impala import vtrace from ray.rllib.agents.impala.vtrace_policy import _make_time_major, \ - BEHAVIOUR_LOGITS, clip_gradients, validate_config, choose_optimizer + clip_gradients, validate_config, choose_optimizer from ray.rllib.evaluation.postprocessing import Postprocessing from ray.rllib.models.tf.tf_action_dist import Categorical from ray.rllib.policy.sample_batch import SampleBatch @@ -244,7 +244,7 @@ def build_appo_surrogate_loss(policy, model, dist_class, train_batch): actions = train_batch[SampleBatch.ACTIONS] dones = train_batch[SampleBatch.DONES] rewards = train_batch[SampleBatch.REWARDS] - behaviour_logits = train_batch[BEHAVIOUR_LOGITS] + behaviour_logits = train_batch[SampleBatch.ACTION_DIST_INPUTS] target_model_out, _ = policy.target_model.from_batch(train_batch) old_policy_behaviour_logits = tf.stop_gradient(target_model_out) @@ -397,8 +397,8 @@ def postprocess_trajectory(policy, return batch -def add_values_and_logits(policy): - out = {BEHAVIOUR_LOGITS: policy.model.last_output()} +def add_values(policy): + out = {} if not policy.config["vtrace"]: out[SampleBatch.VF_PREDS] = policy.model.value_function() return out @@ -446,7 +446,7 @@ AsyncPPOTFPolicy = build_tf_policy( postprocess_fn=postprocess_trajectory, optimizer_fn=choose_optimizer, gradients_fn=clip_gradients, - extra_action_fetches_fn=add_values_and_logits, + extra_action_fetches_fn=add_values, before_init=validate_config, before_loss_init=setup_mixins, after_init=setup_late_mixins, diff --git a/rllib/agents/ppo/ppo_tf_policy.py b/rllib/agents/ppo/ppo_tf_policy.py index e3d643a3e..2afaeb951 100644 --- a/rllib/agents/ppo/ppo_tf_policy.py +++ b/rllib/agents/ppo/ppo_tf_policy.py @@ -1,11 +1,9 @@ import logging import ray -from ray.rllib.agents.impala.vtrace_policy import BEHAVIOUR_LOGITS from ray.rllib.evaluation.postprocessing import compute_advantages, \ Postprocessing from ray.rllib.policy.sample_batch import SampleBatch -from ray.rllib.policy.policy import ACTION_LOGP from ray.rllib.policy.tf_policy import LearningRateSchedule, \ EntropyCoeffSchedule from ray.rllib.policy.tf_policy_template import build_tf_policy @@ -125,8 +123,8 @@ def ppo_surrogate_loss(policy, model, dist_class, train_batch): train_batch[Postprocessing.VALUE_TARGETS], train_batch[Postprocessing.ADVANTAGES], train_batch[SampleBatch.ACTIONS], - train_batch[BEHAVIOUR_LOGITS], - train_batch[ACTION_LOGP], + train_batch[SampleBatch.ACTION_DIST_INPUTS], + train_batch[SampleBatch.ACTION_LOGP], train_batch[SampleBatch.VF_PREDS], action_dist, model.value_function(), @@ -158,11 +156,10 @@ def kl_and_loss_stats(policy, train_batch): } -def vf_preds_and_logits_fetches(policy): - """Adds value function and logits outputs to experience train_batches.""" +def vf_preds_fetches(policy): + """Adds value function outputs to experience train_batches.""" return { SampleBatch.VF_PREDS: policy.model.value_function(), - BEHAVIOUR_LOGITS: policy.model.last_output(), } @@ -270,7 +267,7 @@ PPOTFPolicy = build_tf_policy( get_default_config=lambda: ray.rllib.agents.ppo.ppo.DEFAULT_CONFIG, loss_fn=ppo_surrogate_loss, stats_fn=kl_and_loss_stats, - extra_action_fetches_fn=vf_preds_and_logits_fetches, + extra_action_fetches_fn=vf_preds_fetches, postprocess_fn=postprocess_ppo_gae, gradients_fn=clip_gradients, before_init=setup_config, diff --git a/rllib/agents/ppo/ppo_torch_policy.py b/rllib/agents/ppo/ppo_torch_policy.py index c62ba85ce..7229b251e 100644 --- a/rllib/agents/ppo/ppo_torch_policy.py +++ b/rllib/agents/ppo/ppo_torch_policy.py @@ -1,13 +1,11 @@ import logging import ray -from ray.rllib.agents.impala.vtrace_policy import BEHAVIOUR_LOGITS from ray.rllib.agents.a3c.a3c_torch_policy import apply_grad_clipping from ray.rllib.agents.ppo.ppo_tf_policy import postprocess_ppo_gae, \ setup_config from ray.rllib.evaluation.postprocessing import Postprocessing from ray.rllib.policy.sample_batch import SampleBatch -from ray.rllib.policy.policy import ACTION_LOGP from ray.rllib.policy.torch_policy import EntropyCoeffSchedule, \ LearningRateSchedule from ray.rllib.policy.torch_policy_template import build_torch_policy @@ -128,8 +126,8 @@ def ppo_surrogate_loss(policy, model, dist_class, train_batch): train_batch[Postprocessing.VALUE_TARGETS], train_batch[Postprocessing.ADVANTAGES], train_batch[SampleBatch.ACTIONS], - train_batch[BEHAVIOUR_LOGITS], - train_batch[ACTION_LOGP], + train_batch[SampleBatch.ACTION_DIST_INPUTS], + train_batch[SampleBatch.ACTION_LOGP], train_batch[SampleBatch.VF_PREDS], action_dist, model.value_function(), @@ -162,12 +160,10 @@ def kl_and_loss_stats(policy, train_batch): } -def vf_preds_and_logits_fetches(policy, input_dict, state_batches, model, - action_dist): - """Adds value function and logits outputs to experience train_batches.""" +def vf_preds_fetches(policy, input_dict, state_batches, model, action_dist): + """Adds value function outputs to experience train_batches.""" return { SampleBatch.VF_PREDS: policy.model.value_function(), - BEHAVIOUR_LOGITS: policy.model.last_output(), } @@ -222,7 +218,7 @@ PPOTorchPolicy = build_torch_policy( get_default_config=lambda: ray.rllib.agents.ppo.ppo.DEFAULT_CONFIG, loss_fn=ppo_surrogate_loss, stats_fn=kl_and_loss_stats, - extra_action_out_fn=vf_preds_and_logits_fetches, + extra_action_out_fn=vf_preds_fetches, postprocess_fn=postprocess_ppo_gae, extra_grad_process_fn=apply_grad_clipping, before_init=setup_config, diff --git a/rllib/agents/ppo/tests/test_ppo.py b/rllib/agents/ppo/tests/test_ppo.py index 51a4a644f..7dcb42985 100644 --- a/rllib/agents/ppo/tests/test_ppo.py +++ b/rllib/agents/ppo/tests/test_ppo.py @@ -2,7 +2,6 @@ import numpy as np import unittest import ray -from ray.rllib.agents.impala.vtrace_policy import BEHAVIOUR_LOGITS import ray.rllib.agents.ppo as ppo from ray.rllib.agents.ppo.ppo_tf_policy import postprocess_ppo_gae as \ postprocess_ppo_gae_tf, ppo_surrogate_loss as ppo_surrogate_loss_tf @@ -12,7 +11,6 @@ from ray.rllib.evaluation.postprocessing import Postprocessing from ray.rllib.models.tf.tf_action_dist import Categorical from ray.rllib.models.torch.torch_modelv2 import TorchModelV2 from ray.rllib.models.torch.torch_action_dist import TorchCategorical -from ray.rllib.policy.policy import ACTION_LOGP from ray.rllib.policy.sample_batch import SampleBatch from ray.rllib.utils.framework import try_import_tf from ray.rllib.utils.numpy import fc @@ -109,9 +107,10 @@ class TestPPO(unittest.TestCase): SampleBatch.REWARDS: np.array([1.0, -1.0, .5], dtype=np.float32), SampleBatch.DONES: np.array([False, False, True]), SampleBatch.VF_PREDS: np.array([0.5, 0.6, 0.7], dtype=np.float32), - BEHAVIOUR_LOGITS: np.array( + SampleBatch.ACTION_DIST_INPUTS: np.array( [[-2., 0.5], [-3., -0.3], [-0.1, 2.5]], dtype=np.float32), - ACTION_LOGP: np.array([-0.5, -0.1, -0.2], dtype=np.float32) + SampleBatch.ACTION_LOGP: np.array( + [-0.5, -0.1, -0.2], dtype=np.float32), } for fw in ["tf", "torch"]: @@ -173,17 +172,19 @@ class TestPPO(unittest.TestCase): """ # Calculate expected PPO loss results. dist = dist_class(logits, policy.model) - dist_prev = dist_class(train_batch[BEHAVIOUR_LOGITS], policy.model) + dist_prev = dist_class(train_batch[SampleBatch.ACTION_DIST_INPUTS], + policy.model) expected_logp = dist.logp(train_batch[SampleBatch.ACTIONS]) if isinstance(model, TorchModelV2): expected_rho = np.exp(expected_logp.detach().numpy() - - train_batch.get(ACTION_LOGP)) + train_batch.get(SampleBatch.ACTION_LOGP)) # KL(prev vs current action dist)-loss component. kl = np.mean(dist_prev.kl(dist).detach().numpy()) # Entropy-loss component. entropy = np.mean(dist.entropy().detach().numpy()) else: - expected_rho = np.exp(expected_logp - train_batch[ACTION_LOGP]) + expected_rho = np.exp(expected_logp - + train_batch[SampleBatch.ACTION_LOGP]) # KL(prev vs current action dist)-loss component. kl = np.mean(dist_prev.kl(dist)) # Entropy-loss component. diff --git a/rllib/agents/qmix/qmix_policy.py b/rllib/agents/qmix/qmix_policy.py index 3fdd6a9e4..50febd1bd 100644 --- a/rllib/agents/qmix/qmix_policy.py +++ b/rllib/agents/qmix/qmix_policy.py @@ -216,6 +216,8 @@ class QMixTorchPolicy(Policy): name="target_model", default_model=RNNModel).to(self.device) + self.exploration = self._create_exploration() + # Setup the mixer network. if config["mixer"] is None: self.mixer = None diff --git a/rllib/agents/sac/sac_policy.py b/rllib/agents/sac/sac_policy.py index 09a6a5519..4914a7b89 100644 --- a/rllib/agents/sac/sac_policy.py +++ b/rllib/agents/sac/sac_policy.py @@ -1,19 +1,19 @@ -from gym.spaces import Box, Discrete import logging -import numpy as np +import numpy as np import ray import ray.experimental.tf_utils -from ray.rllib.agents.sac.sac_model import SACModel +from gym.spaces import Box, Discrete from ray.rllib.agents.ddpg.noop_model import NoopModel from ray.rllib.agents.dqn.dqn_policy import postprocess_nstep_and_prio, \ PRIO_WEIGHTS -from ray.rllib.policy.sample_batch import SampleBatch -from ray.rllib.policy.tf_policy import TFPolicy -from ray.rllib.policy.tf_policy_template import build_tf_policy +from ray.rllib.agents.sac.sac_model import SACModel from ray.rllib.models import ModelCatalog from ray.rllib.models.tf.tf_action_dist import (Categorical, SquashedGaussian, DiagGaussian) +from ray.rllib.policy.sample_batch import SampleBatch +from ray.rllib.policy.tf_policy import TFPolicy +from ray.rllib.policy.tf_policy_template import build_tf_policy from ray.rllib.utils import try_import_tf, try_import_tfp from ray.rllib.utils.annotations import override from ray.rllib.utils.error import UnsupportedSpaceException @@ -100,31 +100,21 @@ def get_dist_class(config, action_space): return action_dist_class -def get_log_likelihood(policy, model, actions, input_dict, obs_space, - action_space, config): - model_out, _ = model({ - "obs": input_dict[SampleBatch.CUR_OBS], +def get_distribution_inputs_and_class(policy, + model, + obs_batch, + *, + explore=True, + **kwargs): + # Get base-model output. + model_out, state_out = model({ + "obs": obs_batch, "is_training": policy._get_is_training_placeholder(), }, [], None) + # Get action model output from base-model output. distribution_inputs = model.get_policy_output(model_out) - action_dist_class = get_dist_class(policy.config, action_space) - return action_dist_class(distribution_inputs, model).logp(actions) - - -def build_action_output(policy, model, input_dict, obs_space, action_space, - explore, config, timestep): - model_out, _ = model({ - "obs": input_dict[SampleBatch.CUR_OBS], - "is_training": policy._get_is_training_placeholder(), - }, [], None) - distribution_inputs = model.get_policy_output(model_out) - action_dist_class = get_dist_class(policy.config, action_space) - - policy.output_actions, policy.sampled_action_logp = \ - policy.exploration.get_exploration_action( - distribution_inputs, action_dist_class, model, timestep, explore) - - return policy.output_actions, policy.sampled_action_logp + action_dist_class = get_dist_class(policy.config, policy.action_space) + return distribution_inputs, action_dist_class, state_out def actor_critic_loss(policy, model, _, train_batch): @@ -477,8 +467,7 @@ SACTFPolicy = build_tf_policy( get_default_config=lambda: ray.rllib.agents.sac.sac.DEFAULT_CONFIG, make_model=build_sac_model, postprocess_fn=postprocess_trajectory, - action_sampler_fn=build_action_output, - log_likelihood_fn=get_log_likelihood, + action_distribution_fn=get_distribution_inputs_and_class, loss_fn=actor_critic_loss, stats_fn=stats, gradients_fn=gradients, diff --git a/rllib/contrib/alpha_zero/core/alpha_zero_policy.py b/rllib/contrib/alpha_zero/core/alpha_zero_policy.py index e002f95c8..29916fb20 100644 --- a/rllib/contrib/alpha_zero/core/alpha_zero_policy.py +++ b/rllib/contrib/alpha_zero/core/alpha_zero_policy.py @@ -14,9 +14,12 @@ class AlphaZeroPolicy(TorchPolicy): action_distribution_class, mcts_creator, env_creator, **kwargs): super().__init__( - observation_space, action_space, config, model, loss, - action_distribution_class - ) + observation_space, + action_space, + config, + model=model, + loss=loss, + action_distribution_class=action_distribution_class) # we maintain an env copy in the policy that is used during mcts # simulations self.env_creator = env_creator diff --git a/rllib/contrib/bandits/exploration.py b/rllib/contrib/bandits/exploration.py index 099b56a03..ec05a31de 100644 --- a/rllib/contrib/bandits/exploration.py +++ b/rllib/contrib/bandits/exploration.py @@ -1,6 +1,6 @@ from typing import Union -from ray.rllib.models.modelv2 import ModelV2 +from ray.rllib.models.action_dist import ActionDistribution from ray.rllib.utils.annotations import override from ray.rllib.utils.exploration.exploration import Exploration from ray.rllib.utils.framework import TensorType @@ -9,44 +9,38 @@ from ray.rllib.utils.framework import TensorType class ThompsonSampling(Exploration): @override(Exploration) def get_exploration_action(self, - distribution_inputs: TensorType, - action_dist_class: type, - model: ModelV2, + action_distribution: ActionDistribution, timestep: Union[int, TensorType], explore: bool = True): if self.framework == "torch": - return self._get_torch_exploration_action(distribution_inputs, - explore, model) + return self._get_torch_exploration_action(action_distribution, + explore) else: raise NotImplementedError - def _get_torch_exploration_action(self, distribution_inputs, explore, - model): + def _get_torch_exploration_action(self, action_dist, explore): if explore: - return distribution_inputs.argmax(dim=1), None + return action_dist.inputs.argmax(dim=1), None else: - scores = model.predict(model.current_obs()) + scores = self.model.predict(self.model.current_obs()) return scores.argmax(dim=1), None class UCB(Exploration): @override(Exploration) def get_exploration_action(self, - distribution_inputs: TensorType, - action_dist_class: type, - model: ModelV2, + action_distribution: ActionDistribution, timestep: Union[int, TensorType], explore: bool = True): if self.framework == "torch": - return self._get_torch_exploration_action(distribution_inputs, - explore, model) + return self._get_torch_exploration_action(action_distribution, + explore) else: raise NotImplementedError - def _get_torch_exploration_action(self, distribution_inputs, explore, - model): + def _get_torch_exploration_action(self, action_dist, explore): if explore: - return distribution_inputs.argmax(dim=1), None + return action_dist.inputs.argmax(dim=1), None else: - scores = model.value_function() + scores = self.model.value_function() return scores.argmax(dim=1), None diff --git a/rllib/contrib/maddpg/maddpg_policy.py b/rllib/contrib/maddpg/maddpg_policy.py index fe18fad42..4865db839 100644 --- a/rllib/contrib/maddpg/maddpg_policy.py +++ b/rllib/contrib/maddpg/maddpg_policy.py @@ -73,14 +73,12 @@ class MADDPGTFPolicy(MADDPGPostprocessing, TFPolicy): "Space {} is not supported.".format(space)) obs_space_n = [ - _make_continuous_space(space) - for _, (_, space, _, - _) in sorted(config["multiagent"]["policies"].items()) + _make_continuous_space(space) for _, (_, space, _, _) in + sorted(config["multiagent"]["policies"].items()) ] act_space_n = [ - _make_continuous_space(space) - for _, (_, _, space, - _) in sorted(config["multiagent"]["policies"].items()) + _make_continuous_space(space) for _, (_, _, space, _) in + sorted(config["multiagent"]["policies"].items()) ] # _____ Placeholders @@ -247,7 +245,8 @@ class MADDPGTFPolicy(MADDPGPostprocessing, TFPolicy): obs_input=obs_ph_n[agent_id], sampled_action=act_sampler, loss=actor_loss + critic_loss, - loss_inputs=loss_inputs) + loss_inputs=loss_inputs, + dist_inputs=actor_feature) self.sess.run(tf.global_variables_initializer()) diff --git a/rllib/evaluation/sampler.py b/rllib/evaluation/sampler.py index 9358f289a..f77461592 100644 --- a/rllib/evaluation/sampler.py +++ b/rllib/evaluation/sampler.py @@ -601,8 +601,8 @@ def _do_policy_eval(tf_sess, to_eval, policies, active_episodes): episodes=[active_episodes[t.env_id] for t in eval_data], timestep=policy.global_timestep) if builder: - for k, v in pending_fetches.items(): - eval_results[k] = builder.get(v) + for pid, v in pending_fetches.items(): + eval_results[pid] = builder.get(v) if log_once("compute_actions_result"): logger.info("Outputs of compute_actions():\n\n{}\n".format( @@ -629,7 +629,11 @@ def _process_policy_eval_results(to_eval, eval_results, active_episodes, for policy_id, eval_data in to_eval.items(): rnn_in_cols = _to_column_format([t.rnn_state for t in eval_data]) - actions, rnn_out_cols, pi_info_cols = eval_results[policy_id][:3] + + actions = eval_results[policy_id][0] + rnn_out_cols = eval_results[policy_id][1] + pi_info_cols = eval_results[policy_id][2] + if len(rnn_in_cols) != len(rnn_out_cols): raise ValueError("Length of RNN in did not match RNN out, got: " "{} vs {}".format(rnn_in_cols, rnn_out_cols)) diff --git a/rllib/examples/centralized_critic.py b/rllib/examples/centralized_critic.py index 8679a3ef1..56b25f33e 100644 --- a/rllib/examples/centralized_critic.py +++ b/rllib/examples/centralized_critic.py @@ -17,7 +17,6 @@ import numpy as np from gym.spaces import Discrete from ray import tune -from ray.rllib.agents.impala.vtrace_policy import BEHAVIOUR_LOGITS from ray.rllib.agents.ppo.ppo import PPOTrainer from ray.rllib.agents.ppo.ppo_tf_policy import PPOTFPolicy, KLCoeffMixin, \ PPOLoss @@ -27,7 +26,7 @@ from ray.rllib.examples.twostep_game import TwoStepGame from ray.rllib.models import ModelCatalog from ray.rllib.policy.sample_batch import SampleBatch from ray.rllib.policy.tf_policy import LearningRateSchedule, \ - EntropyCoeffSchedule, ACTION_LOGP + EntropyCoeffSchedule from ray.rllib.models.tf.tf_modelv2 import TFModelV2 from ray.rllib.models.tf.fcnet_v2 import FullyConnectedNetwork from ray.rllib.utils.explained_variance import explained_variance @@ -145,8 +144,8 @@ def loss_with_central_critic(policy, model, dist_class, train_batch): train_batch[Postprocessing.VALUE_TARGETS], train_batch[Postprocessing.ADVANTAGES], train_batch[SampleBatch.ACTIONS], - train_batch[BEHAVIOUR_LOGITS], - train_batch[ACTION_LOGP], + train_batch[SampleBatch.ACTION_DIST_INPUTS], + train_batch[SampleBatch.ACTION_LOGP], train_batch[SampleBatch.VF_PREDS], action_dist, policy.central_value_out, diff --git a/rllib/examples/rock_paper_scissors_multiagent.py b/rllib/examples/rock_paper_scissors_multiagent.py index 22e73a988..c942bbf82 100644 --- a/rllib/examples/rock_paper_scissors_multiagent.py +++ b/rllib/examples/rock_paper_scissors_multiagent.py @@ -18,15 +18,15 @@ from ray.rllib.policy.policy import Policy from ray.rllib.env.multi_agent_env import MultiAgentEnv from ray.rllib.utils import try_import_tf +parser = argparse.ArgumentParser() +parser.add_argument("--stop", type=int, default=1000) + tf = try_import_tf() ROCK = 0 PAPER = 1 SCISSORS = 2 -parser = argparse.ArgumentParser() -parser.add_argument("--stop", type=int, default=400000) - class RockPaperScissorsEnv(MultiAgentEnv): """Two-player environment for rock paper scissors. @@ -82,6 +82,10 @@ class RockPaperScissorsEnv(MultiAgentEnv): class AlwaysSameHeuristic(Policy): """Pick a random move and stick with it for the entire episode.""" + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.exploration = self._create_exploration() + def get_initial_state(self): return [random.choice([ROCK, PAPER, SCISSORS])] @@ -108,6 +112,10 @@ class AlwaysSameHeuristic(Policy): class BeatLastHeuristic(Policy): """Play the move that would beat the last move of the opponent.""" + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.exploration = self._create_exploration() + def compute_actions(self, obs_batch, state_batches=None, @@ -136,13 +144,16 @@ class BeatLastHeuristic(Policy): pass -def run_same_policy(): +def run_same_policy(args): """Use the same policy for both agents (trivial case).""" - tune.run("PG", config={"env": RockPaperScissorsEnv}) + tune.run( + "PG", + stop={"timesteps_total": args.stop}, + config={"env": RockPaperScissorsEnv}) -def run_heuristic_vs_learned(use_lstm=False, trainer="PG"): +def run_heuristic_vs_learned(args, use_lstm=False, trainer="PG"): """Run heuristic policies vs a learned agent. The learned agent should eventually reach a reward of ~5 with @@ -157,7 +168,6 @@ def run_heuristic_vs_learned(use_lstm=False, trainer="PG"): else: return random.choice(["always_same", "beat_last"]) - args = parser.parse_args() tune.run( trainer, stop={"timesteps_total": args.stop}, @@ -186,7 +196,7 @@ def run_heuristic_vs_learned(use_lstm=False, trainer="PG"): }) -def run_with_custom_entropy_loss(): +def run_with_custom_entropy_loss(args): """Example of customizing the loss function of an existing policy. This performs about the same as the default loss does.""" @@ -202,11 +212,16 @@ def run_with_custom_entropy_loss(): loss_fn=entropy_policy_gradient_loss) EntropyLossPG = PGTrainer.with_updates( name="EntropyPG", get_policy_class=lambda _: EntropyPolicy) - run_heuristic_vs_learned(use_lstm=True, trainer=EntropyLossPG) + run_heuristic_vs_learned(args, use_lstm=True, trainer=EntropyLossPG) if __name__ == "__main__": - # run_same_policy() - # run_heuristic_vs_learned(use_lstm=False) - run_heuristic_vs_learned(use_lstm=False) - # run_with_custom_entropy_loss() + args = parser.parse_args() + run_same_policy(args) + print("run_same_policy: ok.") + run_heuristic_vs_learned(args, use_lstm=True) + print("run_heuristic_vs_learned(w/ lstm): ok.") + run_heuristic_vs_learned(args, use_lstm=False) + print("run_heuristic_vs_learned (w/o lstm): ok.") + run_with_custom_entropy_loss(args) + print("run_with_custom_entropy_loss: ok.") diff --git a/rllib/models/tf/tf_action_dist.py b/rllib/models/tf/tf_action_dist.py index 5edc941cd..892ecf029 100644 --- a/rllib/models/tf/tf_action_dist.py +++ b/rllib/models/tf/tf_action_dist.py @@ -338,8 +338,8 @@ class Deterministic(TFActionDistribution): return self.inputs @override(TFActionDistribution) - def sampled_action_logp(self): - return 0.0 + def logp(self, x): + return tf.zeros_like(self.inputs) @override(TFActionDistribution) def _build_sample_op(self): diff --git a/rllib/policy/dynamic_tf_policy.py b/rllib/policy/dynamic_tf_policy.py index 1fce6941f..945e8a10a 100644 --- a/rllib/policy/dynamic_tf_policy.py +++ b/rllib/policy/dynamic_tf_policy.py @@ -48,7 +48,7 @@ class DynamicTFPolicy(TFPolicy): before_loss_init=None, make_model=None, action_sampler_fn=None, - log_likelihood_fn=None, + action_distribution_fn=None, existing_inputs=None, existing_model=None, get_batch_divisibility_req=None, @@ -72,13 +72,18 @@ class DynamicTFPolicy(TFPolicy): All policy variables should be created in this function. If not specified, a default model will be created. action_sampler_fn (Optional[callable]): An optional callable - returning a tuple of action and action prob tensors given - (policy, model, input_dict, obs_space, action_space, config). - If None, a default action distribution will be used. - log_likelihood_fn (Optional[callable]): A callable, - returning a log-likelihood op. - If None, a default class is used and distribution inputs - (for parameterization) will be generated by a model call. + returning a tuple of action and action prob tensors given + (policy, model, input_dict, obs_space, action_space, config). + If None, a default action distribution will be used. + action_distribution_fn (Optional[callable]): A callable returning + distribution inputs (parameters), a dist-class to generate an + action distribution object from, and internal-state outputs + (or an empty list if not applicable). + Note: No Exploration hooks have to be called from within + `action_distribution_fn`. It's should only perform a simple + forward pass through some model. + If None, pass inputs through `self.model()` to get the + distribution inputs. existing_inputs (OrderedDict): When copying a policy, this specifies an existing dict of placeholders to use instead of defining new ones @@ -89,6 +94,8 @@ class DynamicTFPolicy(TFPolicy): obs_include_prev_action_reward (bool): whether to include the previous action and reward in the model input """ + self.observation_space = obs_space + self.action_space = action_space self.config = config self.framework = "tf" self._loss_fn = loss_fn @@ -129,16 +136,17 @@ class DynamicTFPolicy(TFPolicy): self._seq_lens = tf.placeholder( dtype=tf.int32, shape=[None], name="seq_lens") - if action_sampler_fn: + dist_class = dist_inputs = None + if action_sampler_fn or action_distribution_fn: if not make_model: raise ValueError( - "`make_model` is required if `action_sampler_fn` is given") - self.dist_class = None + "`make_model` is required if `action_sampler_fn` OR " + "`action_distribution_fn` is given") else: - self.dist_class, logit_dim = ModelCatalog.get_action_dist( + dist_class, logit_dim = ModelCatalog.get_action_dist( action_space, self.config["model"]) - # Setup model + # Setup self.model. if existing_model: self.model = existing_model elif make_model: @@ -151,6 +159,9 @@ class DynamicTFPolicy(TFPolicy): self.config["model"], framework="tf") + # Create the Exploration object to use for this Policy. + self.exploration = self._create_exploration() + if existing_inputs: self._state_in = [ v for k, v in existing_inputs.items() @@ -164,27 +175,48 @@ class DynamicTFPolicy(TFPolicy): for s in self.model.get_initial_state() ] - model_out, self._state_out = self.model(self._input_dict, - self._state_in, self._seq_lens) - - # Create the Exploration object to use for this Policy. - self.exploration = self._create_exploration(action_space, config) timestep = tf.placeholder(tf.int32, (), name="timestep") - # Setup custom action sampler. + # Fully customized action generation (e.g., custom policy). if action_sampler_fn: sampled_action, sampled_action_logp = action_sampler_fn( - self, self.model, self._input_dict, obs_space, action_space, - explore, config, timestep) - # Create a default action sampler. + self, + self.model, + obs_batch=self._input_dict[SampleBatch.CUR_OBS], + state_batches=self._state_in, + seq_lens=self._seq_lens, + prev_action_batch=self._input_dict[SampleBatch.PREV_ACTIONS], + prev_reward_batch=self._input_dict[SampleBatch.PREV_REWARDS], + explore=explore, + is_training=self._input_dict["is_training"]) else: - # Using an exploration setup. + # Distribution generation is customized, e.g., DQN, DDPG. + if action_distribution_fn: + dist_inputs, dist_class, self._state_out = \ + action_distribution_fn( + self, self.model, + obs_batch=self._input_dict[SampleBatch.CUR_OBS], + state_batches=self._state_in, + seq_lens=self._seq_lens, + prev_action_batch=self._input_dict[ + SampleBatch.PREV_ACTIONS], + prev_reward_batch=self._input_dict[ + SampleBatch.PREV_REWARDS], + explore=explore, + is_training=self._input_dict["is_training"]) + # Default distribution generation behavior: + # Pass through model. E.g., PG, PPO. + else: + dist_inputs, self._state_out = self.model( + self._input_dict, self._state_in, self._seq_lens) + + action_dist = dist_class(dist_inputs, self.model) + + # Using exploration to get final action (e.g. via sampling). sampled_action, sampled_action_logp = \ self.exploration.get_exploration_action( - model_out, - self.dist_class, - self.model, - timestep, + action_distribution=action_dist, + timestep=timestep, explore=explore) # Phase 1 init. @@ -194,18 +226,6 @@ class DynamicTFPolicy(TFPolicy): else: batch_divisibility_req = 1 - # Generate the log-likelihood op. - log_likelihood = None - # From a given function. - if log_likelihood_fn: - log_likelihood = log_likelihood_fn(self, self.model, action_input, - self._input_dict, obs_space, - action_space, config) - # Create default, iff we have a distribution class. - elif self.dist_class is not None: - log_likelihood = self.dist_class(model_out, - self.model).logp(action_input) - super().__init__( obs_space, action_space, @@ -215,7 +235,8 @@ class DynamicTFPolicy(TFPolicy): action_input=action_input, # for logp calculations sampled_action=sampled_action, sampled_action_logp=sampled_action_logp, - log_likelihood=log_likelihood, + dist_inputs=dist_inputs, + dist_class=dist_class, loss=None, # dynamically initialized on run loss_inputs=[], model=self.model, @@ -260,9 +281,8 @@ class DynamicTFPolicy(TFPolicy): existing_inputs[len(self._loss_inputs) + i])) if rnn_inputs: rnn_inputs.append(("seq_lens", existing_inputs[-1])) - input_dict = OrderedDict( - [(k, existing_inputs[i]) - for i, (k, _) in enumerate(self._loss_inputs)] + rnn_inputs) + input_dict = OrderedDict([(k, existing_inputs[i]) for i, ( + k, _) in enumerate(self._loss_inputs)] + rnn_inputs) instance = self.__class__( self.observation_space, self.action_space, diff --git a/rllib/policy/eager_tf_policy.py b/rllib/policy/eager_tf_policy.py index 9c02f9155..2aa034746 100644 --- a/rllib/policy/eager_tf_policy.py +++ b/rllib/policy/eager_tf_policy.py @@ -9,8 +9,7 @@ import numpy as np from ray.util.debug import log_once from ray.rllib.evaluation.episode import _flatten_action from ray.rllib.models.catalog import ModelCatalog -from ray.rllib.policy.policy import Policy, LEARNER_STATS_KEY, ACTION_PROB, \ - ACTION_LOGP +from ray.rllib.policy.policy import Policy, LEARNER_STATS_KEY from ray.rllib.policy.sample_batch import SampleBatch from ray.rllib.utils import add_mixins from ray.rllib.utils.annotations import override @@ -34,10 +33,13 @@ def _convert_to_tf(x): def _convert_to_numpy(x): - if x is None: - return None + def _map(x): + if isinstance(x, tf.Tensor): + return x.numpy() + return x + try: - return tf.nest.map_structure(lambda component: component.numpy(), x) + return tf.nest.map_structure(_map, x) except AttributeError: raise TypeError( ("Object of type {} has no method to convert to numpy.").format( @@ -176,7 +178,7 @@ def build_eager_tf_policy(name, after_init=None, make_model=None, action_sampler_fn=None, - log_likelihood_fn=None, + action_distribution_fn=None, mixins=None, obs_include_prev_action_reward=True, get_batch_divisibility_req=None): @@ -210,11 +212,11 @@ def build_eager_tf_policy(name, self.config = config self.dist_class = None - - if action_sampler_fn: + if action_sampler_fn or action_distribution_fn: if not make_model: - raise ValueError("`make_model` is required if " - "`action_sampler_fn` is given") + raise ValueError( + "`make_model` is required if `action_sampler_fn` OR " + "`action_distribution_fn` is given") else: self.dist_class, logit_dim = ModelCatalog.get_action_dist( action_space, self.config["model"]) @@ -230,12 +232,11 @@ def build_eager_tf_policy(name, config["model"], framework="tf", ) - + self.exploration = self._create_exploration() self._state_in = [ tf.convert_to_tensor(np.array([s])) for s in self.model.get_initial_state() ] - input_dict = { SampleBatch.CUR_OBS: tf.convert_to_tensor( np.array([observation_space.sample()])), @@ -243,7 +244,13 @@ def build_eager_tf_policy(name, [_flatten_action(action_space.sample())]), SampleBatch.PREV_REWARDS: tf.convert_to_tensor([0.]), } - self.model(input_dict, self._state_in, tf.convert_to_tensor([1])) + + if action_distribution_fn: + dist_inputs, self.dist_class, _ = action_distribution_fn( + self, self.model, input_dict[SampleBatch.CUR_OBS]) + else: + self.model(input_dict, self._state_in, + tf.convert_to_tensor([1])) if before_loss_init: before_loss_init(self, observation_space, action_space, config) @@ -313,12 +320,6 @@ def build_eager_tf_policy(name, self._is_training = False self._state_in = state_batches - if tf.executing_eagerly(): - n = len(obs_batch) - else: - n = obs_batch.shape[0] - seq_lens = tf.ones(n, dtype=tf.int32) - input_dict = { SampleBatch.CUR_OBS: tf.convert_to_tensor(obs_batch), "is_training": tf.constant(False), @@ -331,46 +332,59 @@ def build_eager_tf_policy(name, prev_reward_batch), }) - # Custom sampler fn given (which may handle self.exploration). - if action_sampler_fn is not None: - state_out = [] - action, logp = action_sampler_fn( - self, - self.model, - input_dict, - self.observation_space, - self.action_space, - explore, - self.config, - timestep=timestep) # Use Exploration object. - else: - with tf.variable_creator_scope(_disallow_var_creation): - # Call the exploration before_compute_actions hook. - self.exploration.before_compute_actions(timestep=timestep) - - model_out, state_out = self.model(input_dict, - state_batches, seq_lens) - action, logp = self.exploration.get_exploration_action( - model_out, - self.dist_class, + with tf.variable_creator_scope(_disallow_var_creation): + if action_sampler_fn: + dist_class = dist_inputs = None + state_out = [] + actions, logp = self.action_sampler_fn( + self, self.model, + input_dict[SampleBatch.CUR_OBS], + explore=explore, + timestep=timestep) + else: + # Exploration hook before each forward pass. + self.exploration.before_compute_actions( + timestep=timestep, explore=explore) + + if action_distribution_fn: + dist_inputs, dist_class, state_out = \ + action_distribution_fn( + self, self.model, + input_dict[SampleBatch.CUR_OBS], + explore=explore, timestep=timestep) + else: + dist_class = self.dist_class + dist_inputs, state_out = self.model( + input_dict, state_batches, + tf.convert_to_tensor([1])) + + action_dist = dist_class(dist_inputs, self.model) + + # Get the exploration action from the forward results. + actions, logp = self.exploration.get_exploration_action( + action_distribution=action_dist, timestep=timestep, explore=explore) + # Add default and custom fetches. extra_fetches = {} + # Action-logp and action-prob. if logp is not None: - extra_fetches.update({ - ACTION_PROB: tf.exp(logp), - ACTION_LOGP: logp, - }) + extra_fetches[SampleBatch.ACTION_PROB] = tf.exp(logp) + extra_fetches[SampleBatch.ACTION_LOGP] = logp + # Action-dist inputs. + if dist_inputs is not None: + extra_fetches[SampleBatch.ACTION_DIST_INPUTS] = dist_inputs + # Custom extra fetches. if extra_action_fetches_fn: extra_fetches.update(extra_action_fetches_fn(self)) # Increase our global sampling timestep counter by 1. self.global_timestep += 1 - return action, state_out, extra_fetches + return actions, state_out, extra_fetches @override(Policy) def compute_log_likelihoods(self, @@ -379,6 +393,10 @@ def build_eager_tf_policy(name, state_batches=None, prev_action_batch=None, prev_reward_batch=None): + if action_sampler_fn and action_distribution_fn is None: + raise ValueError("Cannot compute log-prob/likelihood w/o an " + "`action_distribution_fn` and a provided " + "`action_sampler_fn`!") seq_lens = tf.ones(len(obs_batch), dtype=tf.int32) input_dict = { @@ -393,11 +411,15 @@ def build_eager_tf_policy(name, prev_reward_batch), }) - # Custom log_likelihood function given. - if log_likelihood_fn: - log_likelihoods = log_likelihood_fn( - self, self.model, actions, input_dict, - self.observation_space, self.action_space, self.config) + # Exploration hook before each forward pass. + self.exploration.before_compute_actions(explore=False) + + # Action dist class and inputs are generated via custom function. + if action_distribution_fn: + dist_inputs, dist_class, _ = action_distribution_fn( + self, self.model, input_dict[SampleBatch.CUR_OBS]) + action_dist = dist_class(dist_inputs, self.model) + log_likelihoods = action_dist.logp(actions) # Default log-likelihood calculation. else: dist_inputs, _ = self.model(input_dict, state_batches, diff --git a/rllib/policy/policy.py b/rllib/policy/policy.py index a3b20b2a6..660cadc6c 100644 --- a/rllib/policy/policy.py +++ b/rllib/policy/policy.py @@ -10,9 +10,6 @@ from ray.rllib.utils.from_config import from_config # `grad_info` dict returned by learn_on_batch() / compute_grads() via this key. LEARNER_STATS_KEY = "learner_stats" -ACTION_PROB = "action_prob" -ACTION_LOGP = "action_logp" - @DeveloperAPI class Policy(metaclass=ABCMeta): @@ -51,10 +48,12 @@ class Policy(metaclass=ABCMeta): self.observation_space = observation_space self.action_space = action_space self.config = config - self.exploration = self._create_exploration(action_space, config) # The global timestep, broadcast down from time to time from the # driver. self.global_timestep = 0 + # The action distribution class to use for action sampling, if any. + # Child classes may set this. + self.dist_class = None @abstractmethod @DeveloperAPI @@ -363,23 +362,26 @@ class Policy(metaclass=ABCMeta): """ raise NotImplementedError - def _create_exploration(self, action_space, config): + def _create_exploration(self): """Creates the Policy's Exploration object. This method only exists b/c some Trainers do not use TfPolicy nor TorchPolicy, but inherit directly from Policy. Others inherit from TfPolicy w/o using DynamicTfPolicy. TODO(sven): unify these cases.""" + if getattr(self, "exploration", None) is not None: + return self.exploration + exploration = from_config( Exploration, - config.get("exploration_config", {"type": "StochasticSampling"}), - action_space=action_space, - num_workers=config.get("num_workers", 0), - worker_index=config.get("worker_index", 0), + self.config.get("exploration_config", + {"type": "StochasticSampling"}), + action_space=self.action_space, + policy_config=self.config, + model=getattr(self, "model", None), + num_workers=self.config.get("num_workers", 0), + worker_index=self.config.get("worker_index", 0), framework=getattr(self, "framework", "tf")) - # If config is further passed around, it'll contain an already - # instantiated object. - config["exploration_config"] = exploration return exploration diff --git a/rllib/policy/sample_batch.py b/rllib/policy/sample_batch.py index 8c98ab350..c299b4b51 100644 --- a/rllib/policy/sample_batch.py +++ b/rllib/policy/sample_batch.py @@ -27,6 +27,11 @@ class SampleBatch: DONES = "dones" INFOS = "infos" + # Extra action fetches keys. + ACTION_DIST_INPUTS = "action_dist_inputs" + ACTION_PROB = "action_prob" + ACTION_LOGP = "action_logp" + # Uniquely identifies an episode EPS_ID = "eps_id" diff --git a/rllib/policy/tests/test_policy.py b/rllib/policy/tests/test_policy.py index 198f9ea67..58c0fc0a0 100644 --- a/rllib/policy/tests/test_policy.py +++ b/rllib/policy/tests/test_policy.py @@ -5,10 +5,12 @@ from ray.rllib.utils.annotations import override class TestPolicy(Policy): + """A dummy Policy that returns a random (batched) int for compute_actions. """ - A dummy Policy that returns a random (batched) int for compute_actions - and implements all other abstract methods of Policy with "pass". - """ + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.exploration = self._create_exploration() @override(Policy) def compute_actions(self, diff --git a/rllib/policy/tf_policy.py b/rllib/policy/tf_policy.py index 19697e331..368f8b990 100644 --- a/rllib/policy/tf_policy.py +++ b/rllib/policy/tf_policy.py @@ -1,13 +1,12 @@ import errno import logging +import numpy as np import os -import numpy as np import ray import ray.experimental.tf_utils from ray.util.debug import log_once -from ray.rllib.policy.policy import Policy, LEARNER_STATS_KEY, \ - ACTION_PROB, ACTION_LOGP +from ray.rllib.policy.policy import Policy, LEARNER_STATS_KEY from ray.rllib.policy.rnn_sequencing import pad_batch_to_sequences_of_same_size from ray.rllib.policy.sample_batch import SampleBatch from ray.rllib.models.modelv2 import ModelV2 @@ -61,6 +60,8 @@ class TFPolicy(Policy): sampled_action_logp=None, action_input=None, log_likelihood=None, + dist_inputs=None, + dist_class=None, state_inputs=None, state_outputs=None, prev_action_input=None, @@ -97,6 +98,10 @@ class TFPolicy(Policy): logp/log-likelihood calculations. log_likelihood (Optional[Tensor]): Tensor to calculate the log_likelihood (given action_input and obs_input). + dist_class (Optional[type): An optional ActionDistribution class + to use for generating a dist object from distribution inputs. + dist_inputs (Optional[Tensor]): Tensor to calculate the + distribution inputs/parameters. state_inputs (list): list of RNN state input Tensors. state_outputs (list): list of RNN state output Tensors. prev_action_input (Tensor): placeholder for previous actions @@ -118,6 +123,7 @@ class TFPolicy(Policy): self.framework = "tf" super().__init__(observation_space, action_space, config) self.model = model + self.exploration = self._create_exploration() self._sess = sess self._obs_input = obs_input self._prev_action_input = prev_action_input @@ -131,6 +137,8 @@ class TFPolicy(Policy): if self._sampled_action_logp is not None else None) self._action_input = action_input # For logp calculations. + self._distr_inputs = dist_inputs + self.dist_class = dist_class self._log_likelihood = log_likelihood self._state_inputs = state_inputs or [] self._state_outputs = state_outputs or [] @@ -162,8 +170,11 @@ class TFPolicy(Policy): raise ValueError( "seq_lens tensor must be given if state inputs are defined") - # Generate the log-likelihood calculator. - self._log_likelihood = log_likelihood + # The log-likelihood calculator op. + self._log_likelihood = None + if self._distr_inputs is not None and self.dist_class is not None: + self._log_likelihood = self.dist_class( + self._distr_inputs, self.model).logp(self._action_input) def variables(self): """Return the list of all savable variables for this policy.""" @@ -253,19 +264,22 @@ class TFPolicy(Policy): timestep=None, **kwargs): explore = explore if explore is not None else self.config["explore"] + timestep = timestep if timestep is not None else self.global_timestep builder = TFRunBuilder(self._sess, "compute_actions") - fetches = self._build_compute_actions( + to_fetch = self._build_compute_actions( builder, obs_batch, - state_batches, - prev_action_batch, - prev_reward_batch, + state_batches=state_batches, + prev_action_batch=prev_action_batch, + prev_reward_batch=prev_reward_batch, explore=explore, - timestep=timestep - if timestep is not None else self.global_timestep) + timestep=timestep) + # Execute session run to get action (and other fetches). - return builder.get(fetches) + fetched = builder.get(to_fetch) + + return fetched @override(Policy) def compute_log_likelihoods(self, @@ -278,8 +292,10 @@ class TFPolicy(Policy): raise ValueError("Cannot compute log-prob/likelihood w/o a " "self._log_likelihood op!") - # Do the forward pass through the model to capture the parameters - # for the action distribution, then do a logp on that distribution. + # Exploration hook before each forward pass. + self.exploration.before_compute_actions( + explore=False, tf_sess=self.get_session()) + builder = TFRunBuilder(self._sess, "compute_log_likelihoods") # Feed actions (for which we want logp values) into graph. builder.add_feed_dict({self._action_input: actions}) @@ -399,13 +415,18 @@ class TFPolicy(Policy): def extra_compute_action_fetches(self): """Extra values to fetch and return from compute_actions(). - By default we only return action probability info (if present). + By default we return action probability/log-likelihood info + and action distribution inputs (if present). """ - ret = {} + extra_fetches = {} + # Action-logp and action-prob. if self._sampled_action_logp is not None: - ret[ACTION_PROB] = self._sampled_action_prob - ret[ACTION_LOGP] = self._sampled_action_logp - return ret + extra_fetches[SampleBatch.ACTION_PROB] = self._sampled_action_prob + extra_fetches[SampleBatch.ACTION_LOGP] = self._sampled_action_logp + # Action-dist inputs. + if self._distr_inputs is not None: + extra_fetches[SampleBatch.ACTION_DIST_INPUTS] = self._distr_inputs + return extra_fetches @DeveloperAPI def extra_compute_grad_feed_dict(self): @@ -520,6 +541,7 @@ class TFPolicy(Policy): def _build_compute_actions(self, builder, obs_batch, + *, state_batches=None, prev_action_batch=None, prev_reward_batch=None, @@ -528,10 +550,11 @@ class TFPolicy(Policy): timestep=None): explore = explore if explore is not None else self.config["explore"] + timestep = timestep if timestep is not None else self.global_timestep # Call the exploration before_compute_actions hook. self.exploration.before_compute_actions( - timestep=self.global_timestep, tf_sess=self.get_session()) + timestep=timestep, explore=explore, tf_sess=self.get_session()) state_batches = state_batches or [] if len(self._state_inputs) != len(state_batches): @@ -602,6 +625,18 @@ class TFPolicy(Policy): return fetches def _get_loss_inputs_dict(self, batch, shuffle): + """Return a feed dict from a batch. + + Arguments: + batch (SampleBatch): batch of data to derive inputs from + shuffle (bool): whether to shuffle batch sequences. Shuffle may + be done in-place. This only makes sense if you're further + applying minibatch SGD after getting the outputs. + + Returns: + feed dict of data + """ + # Get batch ready for RNNs, if applicable. pad_batch_to_sequences_of_same_size( batch, diff --git a/rllib/policy/tf_policy_template.py b/rllib/policy/tf_policy_template.py index 0eefc9da3..02e2b18cd 100644 --- a/rllib/policy/tf_policy_template.py +++ b/rllib/policy/tf_policy_template.py @@ -11,6 +11,7 @@ tf = try_import_tf() @DeveloperAPI def build_tf_policy(name, + *, loss_fn, get_default_config=None, postprocess_fn=None, @@ -26,7 +27,7 @@ def build_tf_policy(name, after_init=None, make_model=None, action_sampler_fn=None, - log_likelihood_fn=None, + action_distribution_fn=None, mixins=None, get_batch_divisibility_req=None, obs_include_prev_action_reward=True): @@ -82,14 +83,12 @@ def build_tf_policy(name, given (policy, obs_space, action_space, config). All policy variables should be created in this function. If not specified, a default model will be created. - action_sampler_fn (Optional[callable]): An optional callable returning - a tuple of action and action prob tensors given - (policy, model, input_dict, obs_space, action_space, config). - If None, a default action distribution will be used. - log_likelihood_fn (Optional[callable]): A callable, - returning a log-likelihood op. - If None, a default class is used and distribution inputs - (for parameterization) will be generated by a model call. + action_sampler_fn (Optional[callable]): A callable returning a sampled + action and its log-likelihood given some (obs and state) inputs. + action_distribution_fn (Optional[callable]): A callable returning + distribution inputs (parameters), a dist-class to generate an + action distribution object from, and internal-state outputs (or an + empty list if not applicable). mixins (list): list of any class mixins for the returned policy class. These mixins will be applied in order and will have higher precedence than the DynamicTFPolicy class @@ -137,7 +136,7 @@ def build_tf_policy(name, before_loss_init=before_loss_init_wrapper, make_model=make_model, action_sampler_fn=action_sampler_fn, - log_likelihood_fn=log_likelihood_fn, + action_distribution_fn=action_distribution_fn, existing_model=existing_model, existing_inputs=existing_inputs, get_batch_divisibility_req=get_batch_divisibility_req, diff --git a/rllib/policy/torch_policy.py b/rllib/policy/torch_policy.py index 9f2925b22..9df12a28a 100644 --- a/rllib/policy/torch_policy.py +++ b/rllib/policy/torch_policy.py @@ -1,8 +1,7 @@ import numpy as np import time -from ray.rllib.policy.policy import Policy, LEARNER_STATS_KEY, ACTION_PROB, \ - ACTION_LOGP +from ray.rllib.policy.policy import Policy, LEARNER_STATS_KEY from ray.rllib.policy.sample_batch import SampleBatch from ray.rllib.policy.rnn_sequencing import pad_batch_to_sequences_of_same_size from ray.rllib.utils.annotations import override, DeveloperAPI @@ -31,9 +30,12 @@ class TorchPolicy(Policy): observation_space, action_space, config, + *, model, loss, action_distribution_class, + action_sampler_fn=None, + action_distribution_fn=None, max_seq_len=20, get_batch_divisibility_req=None): """Build a policy from policy and loss torch modules. @@ -52,6 +54,18 @@ class TorchPolicy(Policy): train_batch) and returns a single scalar loss. action_distribution_class (ActionDistribution): Class for action distribution. + action_sampler_fn (Optional[callable]): A callable returning a + sampled action and its log-likelihood given some (obs and + state) inputs. + action_distribution_fn (Optional[callable]): A callable returning + distribution inputs (parameters), a dist-class to generate an + action distribution object from, and internal-state outputs + (or an empty list if not applicable). + Note: No Exploration hooks have to be called from within + `action_distribution_fn`. It's should only perform a simple + forward pass through some model. + If None, pass inputs through `self.model()` to get the + distribution inputs. max_seq_len (int): Max sequence length for LSTM training. get_batch_divisibility_req (Optional[callable]): Optional callable that returns the divisibility requirement for sample batches. @@ -61,10 +75,14 @@ class TorchPolicy(Policy): self.device = (torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")) self.model = model.to(self.device) + self.exploration = self._create_exploration() self.unwrapped_model = model # used to support DistributedDataParallel self._loss = loss self._optimizer = self.optimizer() + self.dist_class = action_distribution_class + self.action_sampler_fn = action_sampler_fn + self.action_distribution_fn = action_distribution_fn # If set, means we are using distributed allreduce during learning. self.distributed_world_size = None @@ -100,28 +118,51 @@ class TorchPolicy(Policy): input_dict[SampleBatch.PREV_REWARDS] = prev_reward_batch state_batches = [self._convert_to_tensor(s) for s in state_batches] - # Call the exploration before_compute_actions hook. - self.exploration.before_compute_actions(timestep=timestep) + if self.action_sampler_fn: + dist_class = dist_inputs = None + state_out = [] + actions, logp = self.action_sampler_fn( + self, + self.model, + input_dict[SampleBatch.CUR_OBS], + explore=explore, + timestep=timestep) + else: + # Call the exploration before_compute_actions hook. + self.exploration.before_compute_actions(timestep=timestep) + if self.action_distribution_fn: + dist_inputs, dist_class, state_out = \ + self.action_distribution_fn( + self, self.model, input_dict[SampleBatch.CUR_OBS], + explore=explore, timestep=timestep) + else: + dist_class = self.dist_class + dist_inputs, state_out = self.model( + input_dict, state_batches, seq_lens) + action_dist = dist_class(dist_inputs, self.model) + + # Get the exploration action from the forward results. + actions, logp = \ + self.exploration.get_exploration_action( + action_distribution=action_dist, + timestep=timestep, + explore=explore) - model_out = self.model(input_dict, state_batches, seq_lens) - logits, state = model_out - action_dist = None - actions, logp = \ - self.exploration.get_exploration_action( - logits, self.dist_class, self.model, - timestep, explore) input_dict[SampleBatch.ACTIONS] = actions - extra_action_out = self.extra_action_out(input_dict, state_batches, - self.model, action_dist) + # Add default and custom fetches. + extra_fetches = self.extra_action_out(input_dict, state_batches, + self.model, action_dist) + # Action-logp and action-prob. if logp is not None: logp = convert_to_non_torch_type(logp) - extra_action_out.update({ - ACTION_PROB: np.exp(logp), - ACTION_LOGP: logp - }) - return convert_to_non_torch_type((actions, state, - extra_action_out)) + extra_fetches[SampleBatch.ACTION_PROB] = np.exp(logp) + extra_fetches[SampleBatch.ACTION_LOGP] = logp + # Action-dist inputs. + if dist_inputs is not None: + extra_fetches[SampleBatch.ACTION_DIST_INPUTS] = dist_inputs + return convert_to_non_torch_type((actions, state_out, + extra_fetches)) @override(Policy) def compute_log_likelihoods(self, @@ -130,8 +171,13 @@ class TorchPolicy(Policy): state_batches=None, prev_action_batch=None, prev_reward_batch=None): + + if self.action_sampler_fn and self.action_distribution_fn is None: + raise ValueError("Cannot compute log-prob/likelihood w/o an " + "`action_distribution_fn` and a provided " + "`action_sampler_fn`!") + with torch.no_grad(): - seq_lens = torch.ones(len(obs_batch), dtype=torch.int32) input_dict = self._lazy_tensor_dict({ SampleBatch.CUR_OBS: obs_batch, SampleBatch.ACTIONS: actions @@ -140,9 +186,22 @@ class TorchPolicy(Policy): input_dict[SampleBatch.PREV_ACTIONS] = prev_action_batch if prev_reward_batch: input_dict[SampleBatch.PREV_REWARDS] = prev_reward_batch + seq_lens = torch.ones(len(obs_batch), dtype=torch.int32) - parameters, _ = self.model(input_dict, state_batches, seq_lens) - action_dist = self.dist_class(parameters, self.model) + # Exploration hook before each forward pass. + self.exploration.before_compute_actions(explore=False) + + # Action dist class and inputs are generated via custom function. + if self.action_distribution_fn: + dist_inputs, dist_class, _ = self.action_distribution_fn( + self, self.model, input_dict[SampleBatch.CUR_OBS]) + # Default action-dist inputs calculation. + else: + dist_class = self.dist_class + dist_inputs, _ = self.model(input_dict, state_batches, + seq_lens) + + action_dist = dist_class(dist_inputs, self.model) log_likelihoods = action_dist.logp(input_dict[SampleBatch.ACTIONS]) return log_likelihoods diff --git a/rllib/policy/torch_policy_template.py b/rllib/policy/torch_policy_template.py index d908e9aaf..27672118f 100644 --- a/rllib/policy/torch_policy_template.py +++ b/rllib/policy/torch_policy_template.py @@ -12,6 +12,7 @@ torch, _ = try_import_torch() @DeveloperAPI def build_torch_policy(name, + *, loss_fn, get_default_config=None, stats_fn=None, @@ -21,6 +22,8 @@ def build_torch_policy(name, optimizer_fn=None, before_init=None, after_init=None, + action_sampler_fn=None, + action_distribution_fn=None, make_model_and_action_dist=None, mixins=None, get_batch_divisibility_req=None): @@ -46,6 +49,12 @@ def build_torch_policy(name, policy init that takes the same arguments as the policy constructor after_init (func): optional function to run at the end of policy init that takes the same arguments as the policy constructor + action_sampler_fn (Optional[callable]): A callable returning a sampled + action and its log-likelihood given some (obs and state) inputs. + action_distribution_fn (Optional[callable]): A callable returning + distribution inputs (parameters), a dist-class to generate an + action distribution object from, and internal-state outputs (or an + empty list if not applicable). make_model_and_action_dist (func): optional func that takes the same arguments as policy init and returns a tuple of model instance and torch action distribution class. If not specified, the default @@ -73,14 +82,14 @@ def build_torch_policy(name, before_init(self, obs_space, action_space, config) if make_model_and_action_dist: - self.model, self.dist_class = make_model_and_action_dist( + self.model, dist_class = make_model_and_action_dist( self, obs_space, action_space, config) # Make sure, we passed in a correct Model factory. assert isinstance(self.model, TorchModelV2), \ "ERROR: TorchPolicy::make_model_and_action_dist must " \ "return a TorchModelV2 object!" else: - self.dist_class, logit_dim = ModelCatalog.get_action_dist( + dist_class, logit_dim = ModelCatalog.get_action_dist( action_space, self.config["model"], framework="torch") self.model = ModelCatalog.get_model_v2( obs_space=obs_space, @@ -97,7 +106,9 @@ def build_torch_policy(name, config, model=self.model, loss=loss_fn, - action_distribution_class=self.dist_class, + action_distribution_class=dist_class, + action_sampler_fn=action_sampler_fn, + action_distribution_fn=action_distribution_fn, max_seq_len=config["model"]["max_seq_len"], get_batch_divisibility_req=get_batch_divisibility_req, ) diff --git a/rllib/tests/test_checkpoint_restore.py b/rllib/tests/test_checkpoint_restore.py index 7a35a7163..89caf80c7 100644 --- a/rllib/tests/test_checkpoint_restore.py +++ b/rllib/tests/test_checkpoint_restore.py @@ -77,7 +77,7 @@ def ckpt_restore_test(use_object_store, alg_name, failures): alg2 = cls(config=CONFIGS[alg_name], env="CartPole-v0") env = gym.make("CartPole-v0") - for _ in range(2): + for _ in range(1): res = alg1.train() print("current status: " + str(res)) @@ -87,7 +87,7 @@ def ckpt_restore_test(use_object_store, alg_name, failures): else: alg2.restore(alg1.save()) - for _ in range(5): + for _ in range(2): if "DDPG" in alg_name or "SAC" in alg_name: obs = np.clip( np.random.uniform(size=3), @@ -121,7 +121,7 @@ def export_test(alg_name, failures): else: algo = cls(config=CONFIGS[alg_name], env="CartPole-v0") - for _ in range(2): + for _ in range(1): res = algo.train() print("current status: " + str(res)) diff --git a/rllib/utils/exploration/__init__.py b/rllib/utils/exploration/__init__.py index f364a4af8..fbe213fe4 100644 --- a/rllib/utils/exploration/__init__.py +++ b/rllib/utils/exploration/__init__.py @@ -3,6 +3,7 @@ from ray.rllib.utils.exploration.epsilon_greedy import EpsilonGreedy from ray.rllib.utils.exploration.gaussian_noise import GaussianNoise from ray.rllib.utils.exploration.ornstein_uhlenbeck_noise import \ OrnsteinUhlenbeckNoise +from ray.rllib.utils.exploration.parameter_noise import ParameterNoise from ray.rllib.utils.exploration.per_worker_epsilon_greedy import \ PerWorkerEpsilonGreedy from ray.rllib.utils.exploration.per_worker_gaussian_noise import \ @@ -19,6 +20,7 @@ __all__ = [ "EpsilonGreedy", "GaussianNoise", "OrnsteinUhlenbeckNoise", + "ParameterNoise", "PerWorkerEpsilonGreedy", "PerWorkerGaussianNoise", "PerWorkerOrnsteinUhlenbeckNoise", diff --git a/rllib/utils/exploration/epsilon_greedy.py b/rllib/utils/exploration/epsilon_greedy.py index d0195b6f8..4b038ee43 100644 --- a/rllib/utils/exploration/epsilon_greedy.py +++ b/rllib/utils/exploration/epsilon_greedy.py @@ -1,11 +1,12 @@ from typing import Union +from ray.rllib.models.action_dist import ActionDistribution from ray.rllib.utils.annotations import override from ray.rllib.utils.exploration.exploration import Exploration, TensorType from ray.rllib.utils.framework import try_import_tf, try_import_torch, \ get_variable -from ray.rllib.utils.schedules import PiecewiseSchedule -from ray.rllib.models.modelv2 import ModelV2 +from ray.rllib.utils.from_config import from_config +from ray.rllib.utils.schedules import Schedule, PiecewiseSchedule tf = try_import_tf() torch, _ = try_import_torch() @@ -21,33 +22,34 @@ class EpsilonGreedy(Exploration): def __init__(self, action_space, + *, + framework: str, initial_epsilon=1.0, final_epsilon=0.05, epsilon_timesteps=int(1e5), epsilon_schedule=None, - framework="tf", **kwargs): """Create an EpsilonGreedy exploration class. Args: - action_space (Space): The gym action space used by the environment. initial_epsilon (float): The initial epsilon value to use. final_epsilon (float): The final epsilon value to use. epsilon_timesteps (int): The time step after which epsilon should always be `final_epsilon`. epsilon_schedule (Optional[Schedule]): An optional Schedule object to use (instead of constructing one from the given parameters). - framework (Optional[str]): One of None, "tf", "torch". """ assert framework is not None super().__init__( action_space=action_space, framework=framework, **kwargs) - self.epsilon_schedule = epsilon_schedule or PiecewiseSchedule( - endpoints=[(0, initial_epsilon), - (epsilon_timesteps, final_epsilon)], - outside_value=final_epsilon, - framework=self.framework) + self.epsilon_schedule = \ + from_config(Schedule, epsilon_schedule, framework=framework) or \ + PiecewiseSchedule( + endpoints=[ + (0, initial_epsilon), (epsilon_timesteps, final_epsilon)], + outside_value=final_epsilon, + framework=self.framework) # The current timestep value (tf-var or python int). self.last_timestep = get_variable( @@ -55,18 +57,18 @@ class EpsilonGreedy(Exploration): @override(Exploration) def get_exploration_action(self, - distribution_inputs: TensorType, - action_dist_class: type, - model: ModelV2, + *, + action_distribution: ActionDistribution, timestep: Union[int, TensorType], explore: bool = True): + q_values = action_distribution.inputs if self.framework == "tf": - return self._get_tf_exploration_action_op(distribution_inputs, - explore, timestep) + return self._get_tf_exploration_action_op(q_values, explore, + timestep) else: - return self._get_torch_exploration_action(distribution_inputs, - explore, timestep) + return self._get_torch_exploration_action(q_values, explore, + timestep) def _get_tf_exploration_action_op(self, q_values, explore, timestep): """TF method to produce the tf op for an epsilon exploration action. @@ -113,7 +115,7 @@ class EpsilonGreedy(Exploration): """Torch method to produce an epsilon exploration action. Args: - q_values (Tensor): The Q-values coming from some q-model. + q_values (Tensor): The Q-values coming from some Q-model. Returns: torch.Tensor: The exploration-action. diff --git a/rllib/utils/exploration/exploration.py b/rllib/utils/exploration/exploration.py index 1a4fe15a6..2f1442780 100644 --- a/rllib/utils/exploration/exploration.py +++ b/rllib/utils/exploration/exploration.py @@ -3,6 +3,7 @@ from typing import Union from ray.rllib.utils.framework import check_framework, try_import_tf, \ TensorType +from ray.rllib.models.action_dist import ActionDistribution from ray.rllib.models.modelv2 import ModelV2 from ray.rllib.utils.annotations import DeveloperAPI @@ -18,19 +19,21 @@ class Exploration: implemented exploration schema. """ - def __init__(self, - action_space: Space, - num_workers: int, - worker_index: int, - framework: str = "tf"): + def __init__(self, action_space: Space, *, framework: str, + num_workers: int, worker_index: int, policy_config: dict, + model: ModelV2): """ Args: action_space (Space): The action space in which to explore. + framework (str): One of "tf" or "torch". num_workers (int): The overall number of workers used. worker_index (int): The index of the worker using this class. - framework (str): One of "tf" or "torch". + policy_config (dict): The Policy's config dict. + model (ModelV2): The Policy's model. """ self.action_space = action_space + self.policy_config = policy_config + self.model = model self.num_workers = num_workers self.worker_index = worker_index self.framework = check_framework(framework) @@ -54,9 +57,8 @@ class Exploration: @DeveloperAPI def get_exploration_action(self, - distribution_inputs: TensorType, - action_dist_class: type, - model: ModelV2, + *, + action_distribution: ActionDistribution, timestep: Union[int, TensorType], explore: bool = True): """Returns a (possibly) exploratory action and its log-likelihood. @@ -65,12 +67,9 @@ class Exploration: exploratory action. Args: - distribution_inputs (TensorType): The output coming from the model, - ready for parameterizing a distribution - (e.g. q-values or PG-logits). - action_dist_class (class): The action distribution class - to use. - model (ModelV2): The Model object. + action_distribution (ActionDistribution): The instantiated + ActionDistribution object to work with when creating + exploration actions. timestep (int|TensorType): The current sampling time step. It can be a tensor for TF graph mode, otherwise an integer. explore (bool): True: "Normal" exploration behavior. diff --git a/rllib/utils/exploration/gaussian_noise.py b/rllib/utils/exploration/gaussian_noise.py index 4fa6112e7..1a8430dec 100644 --- a/rllib/utils/exploration/gaussian_noise.py +++ b/rllib/utils/exploration/gaussian_noise.py @@ -1,12 +1,13 @@ from typing import Union +from ray.rllib.models.action_dist import ActionDistribution +from ray.rllib.models.modelv2 import ModelV2 from ray.rllib.utils.annotations import override from ray.rllib.utils.exploration.exploration import Exploration from ray.rllib.utils.exploration.random import Random from ray.rllib.utils.framework import try_import_tf, try_import_torch, \ get_variable, TensorType from ray.rllib.utils.schedules.piecewise_schedule import PiecewiseSchedule -from ray.rllib.models.modelv2 import ModelV2 tf = try_import_tf() torch, _ = try_import_torch() @@ -24,18 +25,18 @@ class GaussianNoise(Exploration): def __init__(self, action_space, *, + framework: str, + model: ModelV2, random_timesteps=1000, stddev=0.1, initial_scale=1.0, final_scale=0.02, scale_timesteps=10000, scale_schedule=None, - framework="tf", **kwargs): """Initializes a GaussianNoise Exploration object. Args: - action_space (Space): The gym action space used by the environment. random_timesteps (int): The number of timesteps for which to act completely randomly. Only after this number of timesteps, the `self.scale` annealing process will start (see below). @@ -50,14 +51,14 @@ class GaussianNoise(Exploration): `random_timesteps` steps. scale_schedule (Optional[Schedule]): An optional Schedule object to use (instead of constructing one from the given parameters). - framework (Optional[str]): One of None, "tf", "torch". """ assert framework is not None - super().__init__(action_space, framework=framework, **kwargs) + super().__init__( + action_space, model=model, framework=framework, **kwargs) self.random_timesteps = random_timesteps self.random_exploration = Random( - action_space, framework=self.framework, **kwargs) + action_space, model=self.model, framework=self.framework, **kwargs) self.stddev = stddev # The `scale` annealing schedule. self.scale_schedule = scale_schedule or PiecewiseSchedule( @@ -72,20 +73,17 @@ class GaussianNoise(Exploration): @override(Exploration) def get_exploration_action(self, - distribution_inputs: TensorType, - action_dist_class: type, - model: ModelV2, + *, + action_distribution: ActionDistribution, timestep: Union[int, TensorType], explore: bool = True): # Adds IID Gaussian noise for exploration, TD3-style. - action_dist = action_dist_class(distribution_inputs, model) - if self.framework == "torch": - return self._get_torch_exploration_action(action_dist, explore, - timestep) + return self._get_torch_exploration_action(action_distribution, + explore, timestep) else: - return self._get_tf_exploration_action_op(action_dist, explore, - timestep) + return self._get_tf_exploration_action_op(action_distribution, + explore, timestep) def _get_tf_exploration_action_op(self, action_dist, explore, timestep): ts = timestep if timestep is not None else self.last_timestep diff --git a/rllib/utils/exploration/ornstein_uhlenbeck_noise.py b/rllib/utils/exploration/ornstein_uhlenbeck_noise.py index 2d746c251..7b582eef2 100644 --- a/rllib/utils/exploration/ornstein_uhlenbeck_noise.py +++ b/rllib/utils/exploration/ornstein_uhlenbeck_noise.py @@ -21,6 +21,7 @@ class OrnsteinUhlenbeckNoise(GaussianNoise): def __init__(self, action_space, *, + framework: str, ou_theta=0.15, ou_sigma=0.2, ou_base_scale=0.1, @@ -29,7 +30,6 @@ class OrnsteinUhlenbeckNoise(GaussianNoise): final_scale=0.02, scale_timesteps=10000, scale_schedule=None, - framework="tf", **kwargs): """Initializes an Ornstein-Uhlenbeck Exploration object. @@ -58,13 +58,13 @@ class OrnsteinUhlenbeckNoise(GaussianNoise): """ super().__init__( action_space, + framework=framework, random_timesteps=random_timesteps, initial_scale=initial_scale, final_scale=final_scale, scale_timesteps=scale_timesteps, scale_schedule=scale_schedule, stddev=1.0, # Force `self.stddev` to 1.0. - framework=framework, **kwargs) self.ou_theta = ou_theta self.ou_sigma = ou_sigma diff --git a/rllib/utils/exploration/parameter_noise.py b/rllib/utils/exploration/parameter_noise.py index 0c5d0ff42..fb41b0b41 100644 --- a/rllib/utils/exploration/parameter_noise.py +++ b/rllib/utils/exploration/parameter_noise.py @@ -48,11 +48,12 @@ class ParameterNoise(Exploration): None for auto-detection/setup. """ assert framework is not None - super().__init__(action_space, framework=framework, **kwargs) - - # TODO(sven): Move these to base-Exploration class. - self.policy_config = policy_config, - self.model = model, + super().__init__( + action_space, + policy_config=policy_config, + model=model, + framework=framework, + **kwargs) self.stddev = get_variable( initial_stddev, framework=self.framework, tf_name="stddev") @@ -197,7 +198,7 @@ class ParameterNoise(Exploration): noisy_action_dist = noise_free_action_dist = None # Adjust the stddev depending on the action (pi)-distance. # Also see [1] for details. - distribution = policy.compute_action_distribution( + _, _, fetches = policy.compute_actions( obs_batch=sample_batch[SampleBatch.CUR_OBS], # TODO(sven): What about state-ins and seq-lens? prev_action_batch=sample_batch.get(SampleBatch.PREV_ACTIONS), @@ -205,8 +206,8 @@ class ParameterNoise(Exploration): explore=self.weights_are_currently_noisy) # Categorical case (e.g. DQN). - if isinstance(distribution, Categorical): - action_dist = softmax(distribution.inputs) + if policy.dist_class is Categorical: + action_dist = softmax(fetches[SampleBatch.ACTION_DIST_INPUTS]) else: # TODO(sven): Other action-dist cases. raise NotImplementedError @@ -215,7 +216,7 @@ class ParameterNoise(Exploration): else: noise_free_action_dist = action_dist - distribution = policy.compute_action_distribution( + _, _, fetches = policy.compute_actions( obs_batch=sample_batch[SampleBatch.CUR_OBS], # TODO(sven): What about state-ins and seq-lens? prev_action_batch=sample_batch.get(SampleBatch.PREV_ACTIONS), @@ -223,8 +224,8 @@ class ParameterNoise(Exploration): explore=not self.weights_are_currently_noisy) # Categorical case (e.g. DQN). - if isinstance(distribution, Categorical): - action_dist = softmax(distribution.inputs) + if policy.dist_class is Categorical: + action_dist = softmax(fetches[SampleBatch.ACTION_DIST_INPUTS]) if not self.weights_are_currently_noisy: noisy_action_dist = action_dist @@ -232,7 +233,7 @@ class ParameterNoise(Exploration): noise_free_action_dist = action_dist # Categorical case (e.g. DQN). - if isinstance(distribution, Categorical): + if policy.dist_class is Categorical: # Calculate KL-divergence (DKL(clean||noisy)) according to [2]. # TODO(sven): Allow KL-divergence to be calculated by our # Distribution classes (don't support off-graph/numpy yet). diff --git a/rllib/utils/exploration/per_worker_epsilon_greedy.py b/rllib/utils/exploration/per_worker_epsilon_greedy.py index ed6d07f90..41ada2e21 100644 --- a/rllib/utils/exploration/per_worker_epsilon_greedy.py +++ b/rllib/utils/exploration/per_worker_epsilon_greedy.py @@ -10,7 +10,7 @@ class PerWorkerEpsilonGreedy(EpsilonGreedy): See Ape-X paper. """ - def __init__(self, action_space, *, num_workers, worker_index, framework, + def __init__(self, action_space, *, framework, num_workers, worker_index, **kwargs): """Create a PerWorkerEpsilonGreedy exploration class. diff --git a/rllib/utils/exploration/per_worker_gaussian_noise.py b/rllib/utils/exploration/per_worker_gaussian_noise.py index 72834849d..0bda8f5e1 100644 --- a/rllib/utils/exploration/per_worker_gaussian_noise.py +++ b/rllib/utils/exploration/per_worker_gaussian_noise.py @@ -10,12 +10,7 @@ class PerWorkerGaussianNoise(GaussianNoise): See Ape-X paper. """ - def __init__(self, - action_space, - *, - num_workers, - worker_index, - framework="tf", + def __init__(self, action_space, *, framework, num_workers, worker_index, **kwargs): """ Args: diff --git a/rllib/utils/exploration/per_worker_ornstein_uhlenbeck_noise.py b/rllib/utils/exploration/per_worker_ornstein_uhlenbeck_noise.py index acc5ec019..5ca92d9e8 100644 --- a/rllib/utils/exploration/per_worker_ornstein_uhlenbeck_noise.py +++ b/rllib/utils/exploration/per_worker_ornstein_uhlenbeck_noise.py @@ -11,12 +11,7 @@ class PerWorkerOrnsteinUhlenbeckNoise(OrnsteinUhlenbeckNoise): See Ape-X paper. """ - def __init__(self, - action_space, - *, - num_workers, - worker_index, - framework="tf", + def __init__(self, action_space, *, framework, num_workers, worker_index, **kwargs): """ Args: diff --git a/rllib/utils/exploration/random.py b/rllib/utils/exploration/random.py index dcc2ec126..b772b8843 100644 --- a/rllib/utils/exploration/random.py +++ b/rllib/utils/exploration/random.py @@ -1,12 +1,12 @@ from gym.spaces import Discrete, MultiDiscrete, Tuple from typing import Union +from ray.rllib.models.action_dist import ActionDistribution from ray.rllib.utils.annotations import override from ray.rllib.utils.exploration.exploration import Exploration from ray.rllib.utils.framework import try_import_tf, try_import_torch, \ TensorType from ray.rllib.utils.tuple_actions import TupleActions -from ray.rllib.models.modelv2 import ModelV2 tf = try_import_tf() torch, _ = try_import_torch() @@ -20,7 +20,7 @@ class Random(Exploration): If explore=False, returns the greedy/max-likelihood action. """ - def __init__(self, action_space, *, framework="tf", **kwargs): + def __init__(self, action_space, *, model, framework, **kwargs): """Initialize a Random Exploration object. Args: @@ -28,7 +28,10 @@ class Random(Exploration): framework (Optional[str]): One of None, "tf", "torch". """ super().__init__( - action_space=action_space, framework=framework, **kwargs) + action_space=action_space, + framework=framework, + model=model, + **kwargs) # Determine py_func types, depending on our action-space. if isinstance(self.action_space, (Discrete, MultiDiscrete)) or \ @@ -40,17 +43,17 @@ class Random(Exploration): @override(Exploration) def get_exploration_action(self, - distribution_inputs: TensorType, - action_dist_class: type, - model: ModelV2, + *, + action_distribution: ActionDistribution, timestep: Union[int, TensorType], explore: bool = True): # Instantiate the distribution object. - action_dist = action_dist_class(distribution_inputs, model) if self.framework == "tf": - return self.get_tf_exploration_action_op(action_dist, explore) + return self.get_tf_exploration_action_op(action_distribution, + explore) else: - return self.get_torch_exploration_action(action_dist, explore) + return self.get_torch_exploration_action(action_distribution, + explore) def get_tf_exploration_action_op(self, action_dist, explore): def true_fn(): diff --git a/rllib/utils/exploration/soft_q.py b/rllib/utils/exploration/soft_q.py index 4f40bc5a2..aa3086c54 100644 --- a/rllib/utils/exploration/soft_q.py +++ b/rllib/utils/exploration/soft_q.py @@ -1,6 +1,12 @@ from gym.spaces import Discrete +from typing import Union +from ray.rllib.models.action_dist import ActionDistribution +from ray.rllib.models.tf.tf_action_dist import Categorical +from ray.rllib.models.torch.torch_action_dist import TorchCategorical +from ray.rllib.utils.annotations import override from ray.rllib.utils.exploration.stochastic_sampling import StochasticSampling +from ray.rllib.utils.framework import TensorType class SoftQ(StochasticSampling): @@ -10,23 +16,32 @@ class SoftQ(StochasticSampling): output divided by the temperature. Returns the argmax iff explore=False. """ - def __init__(self, - action_space, - *, - temperature=1.0, - framework="tf", - **kwargs): + def __init__(self, action_space, *, framework, temperature=1.0, **kwargs): """Initializes a SoftQ Exploration object. Args: action_space (Space): The gym action space used by the environment. temperature (Schedule): The temperature to divide model outputs by before creating the Categorical distribution to sample from. - framework (Optional[str]): One of None, "tf", "torch". + framework (str): One of None, "tf", "torch". """ assert isinstance(action_space, Discrete) - super().__init__( - action_space, - static_params=dict(temperature=temperature), - framework=framework, - **kwargs) + super().__init__(action_space, framework=framework, **kwargs) + self.temperature = temperature + + @override(StochasticSampling) + def get_exploration_action(self, + action_distribution: ActionDistribution, + timestep: Union[int, TensorType], + explore: bool = True): + cls = type(action_distribution) + assert cls in [Categorical, TorchCategorical] + # Re-create the action distribution with the correct temperature + # applied. + dist = cls( + action_distribution.inputs, + self.model, + temperature=self.temperature) + # Delegate to super method. + return super().get_exploration_action( + action_distribution=dist, timestep=timestep, explore=explore) diff --git a/rllib/utils/exploration/stochastic_sampling.py b/rllib/utils/exploration/stochastic_sampling.py index 71fc67ad4..95b91b05e 100644 --- a/rllib/utils/exploration/stochastic_sampling.py +++ b/rllib/utils/exploration/stochastic_sampling.py @@ -1,11 +1,12 @@ from typing import Union +from ray.rllib.models.action_dist import ActionDistribution +from ray.rllib.models.modelv2 import ModelV2 from ray.rllib.utils.annotations import override from ray.rllib.utils.exploration.exploration import Exploration from ray.rllib.utils.framework import try_import_tf, try_import_torch, \ TensorType from ray.rllib.utils.tuple_actions import TupleActions -from ray.rllib.models.modelv2 import ModelV2 tf = try_import_tf() torch, _ = try_import_torch() @@ -20,55 +21,30 @@ class StochasticSampling(Exploration): lowering stddev, temperature, etc.. over time. """ - def __init__(self, - action_space, - *, - static_params=None, - time_dependent_params=None, - framework="tf", + def __init__(self, action_space, *, framework: str, model: ModelV2, **kwargs): """Initializes a StochasticSampling Exploration object. Args: action_space (Space): The gym action space used by the environment. - static_params (Optional[dict]): Parameters to be passed as-is into - the action distribution class' constructor. - time_dependent_params (dict): Parameters to be evaluated based on - `timestep` and then passed into the action distribution - class' constructor. - framework (Optional[str]): One of None, "tf", "torch". + framework (str): One of None, "tf", "torch". """ assert framework is not None - super().__init__(action_space, framework=framework, **kwargs) - - self.static_params = static_params or {} - - # TODO(sven): Support scheduled params whose values depend on timestep - # and that will be passed into the distribution's c'tor. - self.time_dependent_params = time_dependent_params or {} + super().__init__( + action_space, model=model, framework=framework, **kwargs) @override(Exploration) def get_exploration_action(self, - distribution_inputs: TensorType, - action_dist_class: type, - model: ModelV2, + *, + action_distribution: ActionDistribution, timestep: Union[int, TensorType], explore: bool = True): - kwargs = self.static_params.copy() - - # TODO(sven): create schedules for these via easy-config patterns - # These can be used anywhere in configs, where schedules are wanted: - # e.g. lr=[0.003, 0.00001, 100k] <- linear anneal from 0.003, to - # 0.00001 over 100k ts. - # if self.time_dependent_params: - # for k, v in self.time_dependent_params: - # kwargs[k] = v(timestep) - action_dist = action_dist_class(distribution_inputs, model, **kwargs) - if self.framework == "torch": - return self._get_torch_exploration_action(action_dist, explore) + return self._get_torch_exploration_action(action_distribution, + explore) else: - return self._get_tf_exploration_action_op(action_dist, explore) + return self._get_tf_exploration_action_op(action_distribution, + explore) def _get_tf_exploration_action_op(self, action_dist, explore): sample = action_dist.sample() diff --git a/rllib/utils/exploration/tests/test_explorations.py b/rllib/utils/exploration/tests/test_explorations.py index 884503545..c5329f112 100644 --- a/rllib/utils/exploration/tests/test_explorations.py +++ b/rllib/utils/exploration/tests/test_explorations.py @@ -35,7 +35,9 @@ def do_test_explorations(run, run in [ddpg.DDPGTrainer, dqn.DQNTrainer, dqn.SimpleQTrainer, impala.ImpalaTrainer, sac.SACTrainer, td3.TD3Trainer]: continue - elif fw == "eager" and run in [ddpg.DDPGTrainer, td3.TD3Trainer]: + elif fw == "eager" and run in [ + ddpg.DDPGTrainer, sac.SACTrainer, td3.TD3Trainer + ]: continue print("Testing {} in framework={}".format(run, fw)) diff --git a/rllib/utils/tests/test_framework_agnostic_components.py b/rllib/utils/tests/test_framework_agnostic_components.py index c7073de61..932f3d33b 100644 --- a/rllib/utils/tests/test_framework_agnostic_components.py +++ b/rllib/utils/tests/test_framework_agnostic_components.py @@ -121,8 +121,11 @@ class TestFrameWorkAgnosticComponents(unittest.TestCase): Exploration, { "type": "EpsilonGreedy", "action_space": Discrete(2), + "framework": "tf", "num_workers": 0, "worker_index": 0, + "policy_config": {}, + "model": None }) check(component.epsilon_schedule.outside_value, 0.05) # default