class_docstr: Linear Bandit Policy to be used by LinUCB, LinTS and possibly others.
class_docstr_tok: ["Linear","Bandit","Policy","to","be","used","by","LinUCB",",","LinTS","and","possibly","others","."]
code: def _distribution(self, time_step, policy_state):
      observation = time_step.observation
      if self.observation_and_action_constraint_splitter is not None:
          (
              observation,
              _,
          ) = self.observation_and_action_constraint_splitter(
              observation
          )
      observation = tf.nest.map_structure(
          lambda o: tf.cast(o, dtype=self._dtype), observation
      )
      global_observation, arm_observations = self._split_observation(
          observation
      )
  
      if self._add_bias:
          # The bias is added via a constant 1 feature.
          global_observation = tf.concat(
              [
                  global_observation,
                  tf.ones(
                      [tf.shape(global_observation)[0], 1],
                      dtype=self._dtype,
                  ),
              ],
              axis=1,
          )
      # Check the shape of the observation matrix. The observations can be
      # batched.
      if not global_observation.shape.is_compatible_with(
          [None, self._global_context_dim]
      ):
          raise ValueError(
              "Global observation shape is expected to be {}. Got {}.".format(
                  [None, self._global_context_dim],
                  global_observation.shape.as_list(),
              )
          )
      global_observation = tf.reshape(
          global_observation, [-1, self._global_context_dim]
      )
      if (self._theta is not None) and (self._alpha < 1e-12):
          # When self._theta is provided and the exploration parameter is nearly
          # zero, inference only needs the predicted mean reward, which can be
          # computed efficiently by using the theta vectors.
          est_rewards = self._predict_mean_reward(
              self._theta, global_observation, arm_observations
          )
          rewards_for_argmax = est_rewards
      else:
          (
              est_rewards,
              est_variances,
          ) = self._predict_mean_reward_and_variance(
              global_observation, arm_observations
          )
          if (
              self._exploration_strategy
              == ExplorationStrategy.optimistic
          ):
              rewards_for_argmax = est_rewards + self._alpha * tf.sqrt(
                  est_variances
              )
          elif (
              self._exploration_strategy == ExplorationStrategy.sampling
          ):
              mu_sampler = tfd.Normal(
                  loc=est_rewards,
                  scale=self._alpha * tf.sqrt(est_variances),
              )
              rewards_for_argmax = mu_sampler.sample()
          else:
              raise ValueError(
                  "Exploraton strategy %s not implemented."
                  % self._exploration_strategy
              )
  
      mask = constraints.construct_mask_from_multiple_sources(
          time_step.observation,
          self._observation_and_action_constraint_splitter,
          (),
          self._num_actions,
      )
      if mask is not None:
          chosen_actions = policy_utilities.masked_argmax(
              rewards_for_argmax,
              mask,
              output_type=tf.nest.flatten(self._action_spec)[0].dtype,
          )
      else:
          chosen_actions = tf.argmax(
              rewards_for_argmax,
              axis=-1,
              output_type=tf.nest.flatten(self._action_spec)[0].dtype,
          )
  
      action_distributions = tfp.distributions.Deterministic(
          loc=chosen_actions
      )
  
      policy_info = policy_utilities.populate_policy_info(
          arm_observations,
          chosen_actions,
          rewards_for_argmax,
          est_rewards,
          self._emit_policy_info,
          self._accepts_per_arm_features,
      )
  
      return policy_step.PolicyStep(
          action_distributions, policy_state, policy_info
      )
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import enum
  from typing import Optional, Sequence, Text, Tuple
  import tensorflow
  import tensorflow_probability
  from tf_agents.bandits.policies import constraints
  from tf_agents.bandits.policies import linalg
  from tf_agents.bandits.specs import utils
  from tf_agents.policies import tf_policy
  from tf_agents.policies import utils
  from tf_agents.specs import tensor_spec
  from tf_agents.trajectories import policy_step
  from tf_agents.typing import types
label_desc: Use image node_id 5 for calling the LinearBanditPolicy obj's underlying member method code with example usage: obj._distribution(time_step, policy_state) and returns: policy_step