code: def testComputeLossWithArmFeatures(self):
      obs_spec = bandit_spec_utils.create_per_arm_observation_spec(
          global_dim=2, per_arm_dim=3, max_num_actions=3
      )
      time_step_spec = ts.time_step_spec(obs_spec)
      constraint_net = global_and_arm_feature_network.create_feed_forward_common_tower_network(
          obs_spec,
          global_layers=(4,),
          arm_layers=(4,),
          common_layers=(4,),
      )
      neural_constraint = constraints.NeuralConstraint(
          time_step_spec,
          self._action_spec,
          constraint_network=constraint_net,
      )
  
      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(18), shape=[2, 3, 3]),
              dtype=tf.float32,
          ),
      }
      actions = tf.constant([0, 1], dtype=tf.int32)
      rewards = tf.constant([0.5, 3.0], dtype=tf.float32)
  
      init_op = neural_constraint.initialize()
      if not tf.executing_eagerly():
          with self.cached_session() as sess:
              common.initialize_uninitialized_variables(sess)
              self.assertIsNone(sess.run(init_op))
      loss = neural_constraint.compute_loss(
          observations, actions, rewards
      )
      self.assertGreater(self.evaluate(loss), 0.0)
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import numpy
  import tensorflow
  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.specs import tensor_spec
  from tf_agents.trajectories import time_step
  from tf_agents.utils import common
label_desc: Use image node_id 6 for calling the NeuralConstraintTest obj's underlying member method code with example usage: obj.testComputeLossWithArmFeatures() without return types