code: def get_simple_q_tf_policy(
      name: str, base: Type[Union[DynamicTFPolicyV2, EagerTFPolicyV2]]
  ) -> Type:
      """Construct a SimpleQTFPolicy inheriting either dynamic or eager base policies.
  
      Args:
          base: Base class for this policy. DynamicTFPolicyV2 or EagerTFPolicyV2.
  
      Returns:
          A TF Policy to be used with MAMLTrainer.
      """
  
      class SimpleQTFPolicy(
          LearningRateSchedule, TargetNetworkMixin, base
      ):
          def __init__(
              self,
              obs_space,
              action_space,
              config,
              existing_model=None,
              existing_inputs=None,
          ):
              # First thing first, enable eager execution if necessary.
              base.enable_eager_execution_if_necessary()
              # Initialize base class.
              base.__init__(
                  self,
                  obs_space,
                  action_space,
                  config,
                  existing_inputs=existing_inputs,
                  existing_model=existing_model,
              )
  
              LearningRateSchedule.__init__(
                  self, config["lr"], config["lr_schedule"]
              )
  
              # Note: this is a bit ugly, but loss and optimizer initialization must
              # happen after all the MixIns are initialized.
              self.maybe_initialize_optimizer_and_loss()
  
              TargetNetworkMixin.__init__(self)
  
          @override(base)
          def make_model(self) -> ModelV2:
              """Builds Q-model and target Q-model for Simple Q learning."""
              model, self.target_model = make_q_models(self)
              return model
  
          @override(base)
          def action_distribution_fn(
              self,
              model: ModelV2,
              *,
              obs_batch: TensorType,
              state_batches: TensorType,
              **kwargs,
          ) -> Tuple[TensorType, type, List[TensorType]]:
              # Compute the Q-values for each possible action, using our Q-value network.
              q_vals = self._compute_q_values(
                  self.model, obs_batch, is_training=False
              )
              return q_vals, Categorical, state_batches
  
          def xyz_compute_actions(
              self,
              *,
              input_dict,
              explore=True,
              timestep=None,
              episodes=None,
              is_training=False,
              **kwargs,
          ) -> Tuple[
              TensorStructType,
              List[TensorType],
              Dict[str, TensorStructType],
          ]:
              if timestep is None:
                  timestep = self.global_timestep
              # Compute the Q-values for each possible action, using our Q-value network.
              q_vals = self._compute_q_values(
                  self.model,
                  input_dict[SampleBatch.OBS],
                  is_training=is_training,
              )
              # Use a Categorical distribution for the exploration component.
              # This way, it may either sample storchastically (e.g. when using SoftQ)
              # or deterministically/greedily (e.g. when using EpsilonGreedy).
              distribution = Categorical(q_vals, self.model)
              # Call the exploration component's `get_exploration_action` method to
              # explore, if necessary.
              actions, logp = self.exploration.get_exploration_action(
                  action_distribution=distribution,
                  timestep=timestep,
                  explore=explore,
              )
              # Return (exploration) actions, state_outs (empty list), and extra outs.
              return (
                  actions,
                  [],
                  {
                      "q_values": q_vals,
                      SampleBatch.ACTION_LOGP: logp,
                      SampleBatch.ACTION_PROB: tf.exp(logp),
                      SampleBatch.ACTION_DIST_INPUTS: q_vals,
                  },
              )
  
          @override(base)
          def loss(
              self,
              model: Union[ModelV2, "tf.keras.Model"],
              dist_class: Type[TFActionDistribution],
              train_batch: SampleBatch,
          ) -> Union[TensorType, List[TensorType]]:
              # q network evaluation
              q_t = self._compute_q_values(
                  self.model, train_batch[SampleBatch.CUR_OBS]
              )
  
              # target q network evalution
              q_tp1 = self._compute_q_values(
                  self.target_model,
                  train_batch[SampleBatch.NEXT_OBS],
              )
              if not hasattr(self, "q_func_vars"):
                  self.q_func_vars = model.variables()
                  self.target_q_func_vars = (
                      self.target_model.variables()
                  )
  
              # q scores for actions which we know were selected in the given state.
              one_hot_selection = tf.one_hot(
                  tf.cast(train_batch[SampleBatch.ACTIONS], tf.int32),
                  self.action_space.n,
              )
              q_t_selected = tf.reduce_sum(q_t * one_hot_selection, 1)
  
              # compute estimate of best possible value starting from state at t + 1
              dones = tf.cast(
                  train_batch[SampleBatch.TERMINATEDS], tf.float32
              )
              q_tp1_best_one_hot_selection = tf.one_hot(
                  tf.argmax(q_tp1, 1), self.action_space.n
              )
              q_tp1_best = tf.reduce_sum(
                  q_tp1 * q_tp1_best_one_hot_selection, 1
              )
              q_tp1_best_masked = (1.0 - dones) * q_tp1_best
  
              # compute RHS of bellman equation
              q_t_selected_target = (
                  tf.cast(train_batch[SampleBatch.REWARDS], tf.float32)
                  + self.config["gamma"] * q_tp1_best_masked
              )
  
              # compute the error (potentially clipped)
              td_error = q_t_selected - tf.stop_gradient(
                  q_t_selected_target
              )
              loss = tf.reduce_mean(huber_loss(td_error))
  
              # save TD error as an attribute for outside access
              self.td_error = td_error
  
              return loss
  
          @override(base)
          def compute_gradients_fn(
              self, optimizer: LocalOptimizer, loss: TensorType
          ) -> ModelGradients:
              return compute_gradients(self, optimizer, loss)
  
          @override(base)
          def extra_learn_fetches_fn(self) -> Dict[str, TensorType]:
              return {
                  "td_error": self.td_error,
                  "learner_stats": {"cur_lr": self.cur_lr},
              }
  
          @override(base)
          def stats_fn(
              self, train_batch: SampleBatch
          ) -> Dict[str, TensorType]:
              """Returns the learning rate in a stats dict.
  
              Args:
                  policy: The Policy object.
                  train_batch: The data used for training.
  
              Returns:
                  Dict[str, TensorType]: The stats dict.
              """
  
              return {
                  "cur_lr": self.cur_lr,
              }
  
          def _compute_q_values(
              self,
              model: ModelV2,
              obs_batch: TensorType,
              is_training=None,
          ) -> TensorType:
              _is_training = (
                  is_training
                  if is_training is not None
                  else self._get_is_training_placeholder()
              )
              model_out, _ = model(
                  SampleBatch(obs=obs_batch, _is_training=_is_training),
                  [],
                  None,
              )
  
              return model_out
  
      SimpleQTFPolicy.__name__ = name
      SimpleQTFPolicy.__qualname__ = name
  
      return SimpleQTFPolicy
  
file_path: ray/rllib_contrib/simple_q/src/rllib_simple_q/simple_q/simple_q_tf_policy.py
imports: import logging
  from typing import Dict, List, Tuple, Type, Union
  from rllib_simple_q.simple_q.utils import make_q_models
  from ray.rllib.models.modelv2 import ModelV2
  from ray.rllib.models.tf.tf_action_dist import Categorical, TFActionDistribution
  from ray.rllib.policy.dynamic_tf_policy_v2 import DynamicTFPolicyV2
  from ray.rllib.policy.eager_tf_policy_v2 import EagerTFPolicyV2
  from ray.rllib.policy.sample_batch import SampleBatch
  from ray.rllib.policy.tf_mixins import LearningRateSchedule, TargetNetworkMixin, compute_gradients
  from ray.rllib.utils.annotations import override
  from ray.rllib.utils.framework import try_import_tf
  from ray.rllib.utils.tf_utils import huber_loss
  from ray.rllib.utils.typing import LocalOptimizer, ModelGradients, TensorStructType, TensorType
label_desc: Use image node_id 1 for calling a global function with example usage: get_simple_q_tf_policy(name, base) and returns: SimpleQTFPolicy, model, q_vals,  Categorical,  state_batches, actions,  list,  dict, loss, compute_gradients, dict, dict, model_out