[RLlib] Exploration API: Policy changes needed for forward pass noisifications. (#7798)

* Rollback.

* WIP.

* WIP.

* LINT.

* WIP.

* Fix.

* Fix.

* Fix.

* LINT.

* Fix (SAC does currently not support eager).

* Fix.

* WIP.

* LINT.

* Update rllib/evaluation/sampler.py

Co-Authored-By: Eric Liang <ekhliang@gmail.com>

* Update rllib/evaluation/sampler.py

Co-Authored-By: Eric Liang <ekhliang@gmail.com>

* Update rllib/utils/exploration/exploration.py

Co-Authored-By: Eric Liang <ekhliang@gmail.com>

* Update rllib/utils/exploration/exploration.py

Co-Authored-By: Eric Liang <ekhliang@gmail.com>

* WIP.

* WIP.

* Fix.

* LINT.

* LINT.

* Fix and LINT.

* WIP.

* WIP.

* WIP.

* WIP.

* Fix.

* LINT.

* Fix.

* Fix and LINT.

* Update rllib/utils/exploration/exploration.py

* Update rllib/policy/dynamic_tf_policy.py

Co-Authored-By: Eric Liang <ekhliang@gmail.com>

* Update rllib/policy/dynamic_tf_policy.py

Co-Authored-By: Eric Liang <ekhliang@gmail.com>

* Update rllib/policy/dynamic_tf_policy.py

Co-Authored-By: Eric Liang <ekhliang@gmail.com>

* Fixes.

* LINT.

* WIP.

Co-authored-by: Eric Liang <ekhliang@gmail.com>

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",
         ]
 )

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())
 

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,

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},

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],

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,

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,

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,

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.

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

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,

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

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

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())
 

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))

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,

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.")

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):

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,

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,

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
 
 

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"
 

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,

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,

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,

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
 

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,
             )

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))
 

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",

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.

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.

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

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

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).

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.
 

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:

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:

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():

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)

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()

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))

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