code: def testTrainPerArmAgent(self):
      obs_spec = bandit_spec_utils.create_per_arm_observation_spec(
          2, 3, 4, add_num_actions_feature=True
      )
      time_step_spec = ts.time_step_spec(
          observation_spec=obs_spec,
          reward_spec=tensor_spec.TensorSpec([3], tf.float32),
      )
      action_spec = tensor_spec.BoundedTensorSpec((), tf.int32, 0, 3)
      networks_and_loss_fns = [
          (
              global_and_arm_feature_network.create_feed_forward_common_tower_network(
                  obs_spec, (4, 3), (3, 4), (4, 2)
              ),
              tf.compat.v1.losses.mean_squared_error,
          )
          for _ in range(3)
      ]
      optimizer = tf.compat.v1.train.GradientDescentOptimizer(
          learning_rate=0.01
      )
      agent = greedy_multi_objective_agent.GreedyMultiObjectiveNeuralAgent(
          time_step_spec,
          action_spec,
          self._scalarizer,
          objective_network_and_loss_fn_sequence=networks_and_loss_fns,
          accepts_per_arm_features=True,
          optimizer=optimizer,
      )
      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,
          ),
          bandit_spec_utils.NUM_ACTIONS_FEATURE_KEY: tf.ones(
              [2], dtype=tf.int32
          ),
      }
      actions = np.array([0, 3], dtype=np.int32)
      objectives = np.array([[1, 2, 3], [6, 5, 4]], dtype=np.float32)
      initial_step, final_step = _get_initial_and_final_steps(
          observations, objectives
      )
      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
      )
      agent.train(experience, None)
      self.evaluate(tf.compat.v1.initialize_all_variables())
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import collections
  from typing import List, Tuple
  import numpy
  import tensorflow
  from tf_agents.bandits.agents import greedy_multi_objective_neural_agent
  from tf_agents.bandits.drivers import driver_utils
  from tf_agents.bandits.multi_objective import multi_objective_scalarizer
  from tf_agents.bandits.networks import global_and_arm_feature_network
  from tf_agents.bandits.networks import heteroscedastic_q_network
  from tf_agents.bandits.specs import utils
  from tf_agents.networks import network
  from tf_agents.policies import utils
  from tf_agents.specs import array_spec
  from tf_agents.specs import tensor_spec
  from tf_agents.trajectories import policy_step
  from tf_agents.trajectories import time_step
  from tf_agents.typing import types
  from tf_agents.utils import common
label_desc: Use image node_id 18 for calling the MultiObjectiveAgentTest obj's underlying member method code with example usage: obj.testTrainPerArmAgent() without return types