Custom action distributions

python/ray/rllib/agents/a3c/a3c_torch_policy.py

@@ -18,7 +18,7 @@ def actor_critic_loss(policy, batch_tensors):
         SampleBatch.CUR_OBS: batch_tensors[SampleBatch.CUR_OBS]
     })  # TODO(ekl) seq lens shouldn't be None
     values = policy.model.value_function()
-    dist = policy.dist_class(logits)
+    dist = policy.dist_class(logits, policy.config["model"])
     log_probs = dist.logp(batch_tensors[SampleBatch.ACTIONS])
     policy.entropy = dist.entropy().mean()
     policy.pi_err = -batch_tensors[Postprocessing.ADVANTAGES].dot(

python/ray/rllib/agents/ars/policies.py

@@ -81,7 +81,7 @@ def __init__(self,
         model = ModelCatalog.get_model({
             "obs": self.inputs
         }, obs_space, action_space, dist_dim, model_config)
-        dist = dist_class(model.outputs)
+        dist = dist_class(model.outputs, model_config=model_config)
         self.sampler = dist.sample()
 
         self.variables = ray.experimental.tf_utils.TensorFlowVariables(

python/ray/rllib/agents/dqn/dqn_policy.py

@@ -107,9 +107,10 @@ def __init__(self,
 
 class QValuePolicy(object):
     def __init__(self, q_values, observations, num_actions, stochastic, eps,
-                 softmax, softmax_temp):
+                 softmax, softmax_temp, model_config):
         if softmax:
-            action_dist = Categorical(q_values / softmax_temp)
+            action_dist = Categorical(
+                q_values / softmax_temp, model_config=model_config)
             self.action = action_dist.sample()
             self.action_prob = action_dist.sampled_action_prob()
             return
@@ -255,7 +256,8 @@ def build_q_networks(policy, q_model, input_dict, obs_space, action_space,
     # Action outputs
     qvp = QValuePolicy(q_values, input_dict[SampleBatch.CUR_OBS],
                        action_space.n, policy.stochastic, policy.eps,
-                       config["soft_q"], config["softmax_temp"])
+                       config["soft_q"], config["softmax_temp"],
+                       config["model"])
     policy.output_actions, policy.action_prob = qvp.action, qvp.action_prob
 
     return policy.output_actions, policy.action_prob

python/ray/rllib/agents/es/policies.py

@@ -59,7 +59,7 @@ def __init__(self, sess, action_space, obs_space, preprocessor,
         model = ModelCatalog.get_model({
             "obs": self.inputs
         }, obs_space, action_space, dist_dim, model_options)
-        dist = dist_class(model.outputs)
+        dist = dist_class(model.outputs, model_config=model_options)
         self.sampler = dist.sample()
 
         self.variables = ray.experimental.tf_utils.TensorFlowVariables(

python/ray/rllib/agents/impala/vtrace.py

@@ -49,13 +49,14 @@
 
 def log_probs_from_logits_and_actions(policy_logits,
                                       actions,
+                                      config,
                                       dist_class=Categorical):
     return multi_log_probs_from_logits_and_actions([policy_logits], [actions],
-                                                   dist_class)[0]
+                                                   dist_class, config)[0]
 
 
-def multi_log_probs_from_logits_and_actions(policy_logits, actions,
-                                            dist_class):
+def multi_log_probs_from_logits_and_actions(policy_logits, actions, dist_class,
+                                            config):
     """Computes action log-probs from policy logits and actions.
 
   In the notation used throughout documentation and comments, T refers to the
@@ -76,6 +77,8 @@ def multi_log_probs_from_logits_and_actions(policy_logits, actions,
       ...,
       [T, B, ...]
       with actions.
+    dist_class: Python class of the action distribution
+    config: Trainer config dict
 
   Returns:
     A list with length of ACTION_SPACE of float32
@@ -97,7 +100,8 @@ def multi_log_probs_from_logits_and_actions(policy_logits, actions,
                                   tf.concat([[-1], a_shape[2:]], axis=0))
         log_probs.append(
             tf.reshape(
-                dist_class(policy_logits_flat).logp(actions_flat),
+                dist_class(policy_logits_flat,
+                           model_config=config["model"]).logp(actions_flat),
                 a_shape[:2]))
 
     return log_probs
@@ -110,6 +114,7 @@ def from_logits(behaviour_policy_logits,
                 rewards,
                 values,
                 bootstrap_value,
+                config,
                 dist_class=Categorical,
                 clip_rho_threshold=1.0,
                 clip_pg_rho_threshold=1.0,
@@ -122,6 +127,7 @@ def from_logits(behaviour_policy_logits,
         rewards,
         values,
         bootstrap_value,
+        config,
         dist_class,
         clip_rho_threshold=clip_rho_threshold,
         clip_pg_rho_threshold=clip_pg_rho_threshold,
@@ -145,6 +151,7 @@ def multi_from_logits(behaviour_policy_logits,
                       rewards,
                       values,
                       bootstrap_value,
+                      config,
                       dist_class,
                       clip_rho_threshold=1.0,
                       clip_pg_rho_threshold=1.0,
@@ -235,9 +242,9 @@ def multi_from_logits(behaviour_policy_logits,
                 discounts, rewards, values, bootstrap_value
             ]):
         target_action_log_probs = multi_log_probs_from_logits_and_actions(
-            target_policy_logits, actions, dist_class)
+            target_policy_logits, actions, dist_class, config)
         behaviour_action_log_probs = multi_log_probs_from_logits_and_actions(
-            behaviour_policy_logits, actions, dist_class)
+            behaviour_policy_logits, actions, dist_class, config)
 
         log_rhos = get_log_rhos(target_action_log_probs,
                                 behaviour_action_log_probs)

python/ray/rllib/agents/impala/vtrace_policy.py

@@ -41,6 +41,7 @@ def __init__(self,
                  bootstrap_value,
                  dist_class,
                  valid_mask,
+                 config,
                  vf_loss_coeff=0.5,
                  entropy_coeff=0.01,
                  clip_rho_threshold=1.0,
@@ -72,6 +73,7 @@ def __init__(self,
             bootstrap_value: A float32 tensor of shape [B].
             dist_class: action distribution class for logits.
             valid_mask: A bool tensor of valid RNN input elements (#2992).
+            config: Trainer config dict.
         """
 
         # Compute vtrace on the CPU for better perf.
@@ -87,7 +89,8 @@ def __init__(self,
                 dist_class=dist_class,
                 clip_rho_threshold=tf.cast(clip_rho_threshold, tf.float32),
                 clip_pg_rho_threshold=tf.cast(clip_pg_rho_threshold,
-                                              tf.float32))
+                                              tf.float32),
+                config=config)
             self.value_targets = self.vtrace_returns.vs
 
         # The policy gradients loss
@@ -196,6 +199,7 @@ def make_time_major(*args, **kw):
         bootstrap_value=make_time_major(values)[-1],
         dist_class=Categorical if is_multidiscrete else policy.dist_class,
         valid_mask=make_time_major(mask, drop_last=True),
+        config=policy.config,
         vf_loss_coeff=policy.config["vf_loss_coeff"],
         entropy_coeff=policy.entropy_coeff,
         clip_rho_threshold=policy.config["vtrace_clip_rho_threshold"],

python/ray/rllib/agents/impala/vtrace_test.py

@@ -98,7 +98,7 @@ def test_log_probs_from_logits_and_actions(self, batch_size):
             0, num_actions - 1, size=(seq_len, batch_size), dtype=np.int32)
 
         action_log_probs_tensor = vtrace.log_probs_from_logits_and_actions(
-            policy_logits, actions)
+            policy_logits, actions, {"model": None})  # dummy config dict
 
         # Ground Truth
         # Using broadcasting to create a mask that indexes action logits
@@ -159,6 +159,8 @@ def test_vtrace_from_logits(self, batch_size):
         clip_rho_threshold = None  # No clipping.
         clip_pg_rho_threshold = None  # No clipping.
 
+        dummy_config = {"model": None}
+
         # Intentionally leaving shapes unspecified to test if V-trace can
         # deal with that.
         placeholders = {
@@ -178,12 +180,15 @@ def test_vtrace_from_logits(self, batch_size):
         from_logits_output = vtrace.from_logits(
             clip_rho_threshold=clip_rho_threshold,
             clip_pg_rho_threshold=clip_pg_rho_threshold,
+            config=dummy_config,
             **placeholders)
 
         target_log_probs = vtrace.log_probs_from_logits_and_actions(
-            placeholders["target_policy_logits"], placeholders["actions"])
+            placeholders["target_policy_logits"], placeholders["actions"],
+            dummy_config)
         behaviour_log_probs = vtrace.log_probs_from_logits_and_actions(
-            placeholders["behaviour_policy_logits"], placeholders["actions"])
+            placeholders["behaviour_policy_logits"], placeholders["actions"],
+            dummy_config)
         log_rhos = target_log_probs - behaviour_log_probs
         ground_truth = (log_rhos, behaviour_log_probs, target_log_probs)
 

python/ray/rllib/agents/marwil/marwil_policy.py

@@ -29,7 +29,7 @@ def __init__(self, state_values, cumulative_rewards):
 
 class ReweightedImitationLoss(object):
     def __init__(self, state_values, cumulative_rewards, logits, actions,
-                 action_space, beta):
+                 action_space, beta, model_config):
         ma_adv_norm = tf.get_variable(
             name="moving_average_of_advantage_norm",
             dtype=tf.float32,
@@ -48,8 +48,8 @@ def __init__(self, state_values, cumulative_rewards, logits, actions,
                 beta * tf.divide(adv, 1e-8 + tf.sqrt(ma_adv_norm)))
 
         # log\pi_\theta(a|s)
-        dist_cls, _ = ModelCatalog.get_action_dist(action_space, {})
-        action_dist = dist_cls(logits)
+        dist_cls, _ = ModelCatalog.get_action_dist(action_space, model_config)
+        action_dist = dist_cls(logits, model_config=model_config)
         logprobs = action_dist.logp(actions)
 
         self.loss = -1.0 * tf.reduce_mean(
@@ -106,7 +106,7 @@ def __init__(self, observation_space, action_space, config):
             self.p_func_vars = scope_vars(scope.name)
 
         # Action outputs
-        action_dist = dist_cls(logits)
+        action_dist = dist_cls(logits, model_config=self.config["model"])
         self.output_actions = action_dist.sample()
 
         # Training inputs
@@ -164,7 +164,8 @@ def _build_value_loss(self, state_values, cum_rwds):
     def _build_policy_loss(self, state_values, cum_rwds, logits, actions,
                            action_space):
         return ReweightedImitationLoss(state_values, cum_rwds, logits, actions,
-                                       action_space, self.config["beta"])
+                                       action_space, self.config["beta"],
+                                       self.config["model"])
 
     @override(TFPolicy)
     def extra_compute_grad_fetches(self):

python/ray/rllib/agents/pg/torch_pg_policy.py

@@ -13,7 +13,8 @@ def pg_torch_loss(policy, batch_tensors):
     logits, _ = policy.model({
         SampleBatch.CUR_OBS: batch_tensors[SampleBatch.CUR_OBS]
     })
-    action_dist = policy.dist_class(logits)
+    action_dist = policy.dist_class(
+        logits, model_config=policy.config["model"])
     log_probs = action_dist.logp(batch_tensors[SampleBatch.ACTIONS])
     # save the error in the policy object
     policy.pi_err = -batch_tensors[Postprocessing.ADVANTAGES].dot(

python/ray/rllib/agents/ppo/appo_policy.py

@@ -112,6 +112,7 @@ def __init__(self,
                  rewards,
                  values,
                  bootstrap_value,
+                 config,
                  dist_class,
                  valid_mask,
                  vf_loss_coeff=0.5,
@@ -143,6 +144,7 @@ def __init__(self,
             rewards: A float32 tensor of shape [T, B].
             values: A float32 tensor of shape [T, B].
             bootstrap_value: A float32 tensor of shape [B].
+            config: Trainer config dict.
             dist_class: action distribution class for logits.
             valid_mask: A bool tensor of valid RNN input elements (#2992).
             vf_loss_coeff (float): Coefficient of the value function loss.
@@ -165,6 +167,7 @@ def reduce_mean_valid(t):
                 rewards=rewards,
                 values=values,
                 bootstrap_value=bootstrap_value,
+                config=config,
                 dist_class=dist_class,
                 clip_rho_threshold=tf.cast(clip_rho_threshold, tf.float32),
                 clip_pg_rho_threshold=tf.cast(clip_pg_rho_threshold,
@@ -251,8 +254,10 @@ def make_time_major(*args, **kw):
         old_policy_behaviour_logits, output_hidden_shape, axis=1)
     unpacked_outputs = tf.split(policy.model_out, output_hidden_shape, axis=1)
     action_dist = policy.action_dist
-    old_policy_action_dist = policy.dist_class(old_policy_behaviour_logits)
-    prev_action_dist = policy.dist_class(behaviour_logits)
+    old_policy_action_dist = policy.dist_class(
+        old_policy_behaviour_logits, model_config=policy.config["model"])
+    prev_action_dist = policy.dist_class(
+        behaviour_logits, model_config=policy.config["model"])
     values = policy.value_function
 
     policy.model_vars = policy.model.variables()
@@ -298,6 +303,7 @@ def make_time_major(*args, **kw):
             rewards=make_time_major(rewards, drop_last=True),
             values=make_time_major(values, drop_last=True),
             bootstrap_value=make_time_major(values)[-1],
+            config=policy.config,
             dist_class=Categorical if is_multidiscrete else policy.dist_class,
             valid_mask=make_time_major(mask, drop_last=True),
             vf_loss_coeff=policy.config["vf_loss_coeff"],

python/ray/rllib/agents/ppo/ppo_policy.py

@@ -39,7 +39,8 @@ def __init__(self,
                  clip_param=0.1,
                  vf_clip_param=0.1,
                  vf_loss_coeff=1.0,
-                 use_gae=True):
+                 use_gae=True,
+                 model_config=None):
         """Constructs the loss for Proximal Policy Objective.
 
         Arguments:
@@ -65,13 +66,15 @@ def __init__(self,
             vf_clip_param (float): Clip parameter for the value function
             vf_loss_coeff (float): Coefficient of the value function loss
             use_gae (bool): If true, use the Generalized Advantage Estimator.
+            model_config (dict): (Optional) model config for use in specifying
+                action distributions.
         """
 
         def reduce_mean_valid(t):
             return tf.reduce_mean(tf.boolean_mask(t, valid_mask))
 
-        dist_cls, _ = ModelCatalog.get_action_dist(action_space, {})
-        prev_dist = dist_cls(logits)
+        dist_cls, _ = ModelCatalog.get_action_dist(action_space, model_config)
+        prev_dist = dist_cls(logits, model_config=model_config)
         # Make loss functions.
         logp_ratio = tf.exp(
             curr_action_dist.logp(actions) - prev_dist.logp(actions))
@@ -129,7 +132,8 @@ def ppo_surrogate_loss(policy, batch_tensors):
         clip_param=policy.config["clip_param"],
         vf_clip_param=policy.config["vf_clip_param"],
         vf_loss_coeff=policy.config["vf_loss_coeff"],
-        use_gae=policy.config["use_gae"])
+        use_gae=policy.config["use_gae"],
+        model_config=policy.config["model"])
 
     return policy.loss_obj.loss
 

python/ray/rllib/agents/ppo/test/test.py

@@ -20,7 +20,7 @@ def testCategorical(self):
         logits = tf.placeholder(tf.float32, shape=(None, 10))
         z = 8 * (np.random.rand(10) - 0.5)
         data = np.tile(z, (num_samples, 1))
-        c = Categorical(logits)
+        c = Categorical(logits, {})  # dummy config dict
         sample_op = c.sample()
         sess = tf.Session()
         sess.run(tf.global_variables_initializer())

python/ray/rllib/examples/custom_loss.py

@@ -67,7 +67,7 @@ def custom_loss(self, policy_loss, loss_inputs):
         print("FYI: You can also use these tensors: {}, ".format(loss_inputs))
 
         # compute the IL loss
-        action_dist = Categorical(logits)
+        action_dist = Categorical(logits, self.options)
         self.policy_loss = policy_loss
         self.imitation_loss = tf.reduce_mean(
             -action_dist.logp(input_ops["actions"]))

python/ray/rllib/examples/custom_torch_policy.py

@@ -18,7 +18,7 @@ def policy_gradient_loss(policy, batch_tensors):
     logits, _ = policy.model({
         SampleBatch.CUR_OBS: batch_tensors[SampleBatch.CUR_OBS]
     })
-    action_dist = policy.dist_class(logits)
+    action_dist = policy.dist_class(logits, policy.config["model"])
     log_probs = action_dist.logp(batch_tensors[SampleBatch.ACTIONS])
     return -batch_tensors[SampleBatch.REWARDS].dot(log_probs)
 

python/ray/rllib/models/action_dist.py

@@ -11,11 +11,14 @@ class ActionDistribution(object):
 
     Args:
       inputs (Tensor): The input vector to compute samples from.
+      model_config (dict): Optional model config dict
+          (as defined in catalog.py)
     """
 
     @DeveloperAPI
-    def __init__(self, inputs):
+    def __init__(self, inputs, model_config):
         self.inputs = inputs
+        self.model_config = model_config
 
     @DeveloperAPI
     def sample(self):
@@ -52,3 +55,22 @@ def multi_entropy(self):
         MultiDiscrete. TODO(ekl) consider removing this.
         """
         return self.entropy()
+
+    @DeveloperAPI
+    @staticmethod
+    def required_model_output_shape(action_space, model_config):
+        """Returns the required shape of an input parameter tensor for a
+        particular action space and an optional dict of distribution-specific
+        options.
+
+        Args:
+            action_space (gym.Space): The action space this distribution will
+                be used for, whose shape attributes will be used to determine
+                the required shape of the input parameter tensor.
+            model_config (dict): Model's config dict (as defined in catalog.py)
+
+        Returns:
+            model_output_shape (int or np.ndarray of ints): size of the
+                required input vector (minus leading batch dimension).
+        """
+        raise NotImplementedError