code: def testTrainPerArmAgentWithConstraint(
      self, agent_class, use_num_samples_list
  ):
      obs_spec = bandit_spec_utils.create_per_arm_observation_spec(
          2, 3, 4
      )
      reward_spec = {
          "reward": tensor_spec.TensorSpec(
              shape=(), dtype=tf.float32, name="reward"
          ),
          "constraint": tensor_spec.TensorSpec(
              shape=(1,), dtype=tf.float32, name="constraint"
          ),
      }
      time_step_spec = ts.time_step_spec(obs_spec, reward_spec)
      reward_net = global_and_arm_feature_network.create_feed_forward_common_tower_network(
          obs_spec, (4, 3), (3, 4), (4, 2)
      )
      optimizer = tf.compat.v1.train.GradientDescentOptimizer(
          learning_rate=0.1
      )
      constraint_net = global_and_arm_feature_network.create_feed_forward_common_tower_network(
          obs_spec, (4, 3), (3, 4), (4, 2)
      )
      neural_constraint = constraints.NeuralConstraint(
          time_step_spec,
          self._action_spec,
          constraint_network=constraint_net,
      )
  
      if use_num_samples_list:
          agent_class = functools.partial(
              agent_class,
              num_samples_list=self._per_arm_num_samples_list,
          )
  
      agent = agent_class(
          time_step_spec,
          self._action_spec,
          reward_network=reward_net,
          accepts_per_arm_features=True,
          optimizer=optimizer,
          constraints=[neural_constraint],
      )
      observations = {
          bandit_spec_utils.GLOBAL_FEATURE_KEY: tf.constant(
              [[1, 2], [3, 4]], dtype=tf.float32
          ),
          bandit_spec_utils.PER_ARM_FEATURE_KEY: tf.cast(
              tf.reshape(tf.range(24), shape=[2, 4, 3]),
              dtype=tf.float32,
          ),
      }
      actions = np.array([0, 3], dtype=np.int32)
      rewards = {
          "reward": np.array([0.5, 3.0], dtype=np.float32),
          "constraint": np.array([[6.0], [4.0]], dtype=np.float32),
      }
      (
          initial_step,
          final_step,
      ) = _get_initial_and_final_steps_nested_rewards(
          observations, rewards
      )
      action_step = policy_step.PolicyStep(
          action=tf.convert_to_tensor(actions),
          info=policy_utilities.PerArmPolicyInfo(
              chosen_arm_features=np.array(
                  [[1, 2, 3], [3, 2, 1]], dtype=np.float32
              )
          ),
      )
      experience = _get_experience(
          initial_step, action_step, final_step
      )
      self.evaluate(tf.compat.v1.initialize_all_variables())
      self.evaluate(agent.train(experience, None).loss)
      if use_num_samples_list:
          self.evaluate(agent.train(experience, None).loss)
          self.assertEqual(
              self.evaluate(
                  self._per_arm_num_samples_list[0].read_value()
              ),
              4,
          )
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import functools
  from absl.testing import parameterized
  import numpy
  import tensorflow
  from tf_agents.bandits.agents import greedy_reward_prediction_agent
  from tf_agents.bandits.agents import neural_boltzmann_agent
  from tf_agents.bandits.agents import neural_epsilon_greedy_agent
  from tf_agents.bandits.agents import neural_falcon_agent
  from tf_agents.bandits.drivers import driver_utils
  from tf_agents.bandits.networks import global_and_arm_feature_network
  from tf_agents.bandits.policies import constraints
  from tf_agents.bandits.specs import utils
  from tf_agents.networks import network
  from tf_agents.policies import utils
  from tf_agents.specs import tensor_spec
  from tf_agents.trajectories import policy_step
  from tf_agents.trajectories import time_step
  from tf_agents.utils import common
label_desc: Use image node_id 15 for calling the AgentTest obj's underlying member method code with example usage: obj.testTrainPerArmAgentWithConstraint(agent_class, use_num_samples_list) without return types