Make average return metric work even if the reward is a vector.

In that case the return is the sum of the reward vector.

PiperOrigin-RevId: 471044125
Change-Id: Ifdb1d4e477fe9bbdc5c2946aaeaa78d7d302b154

Author: bartokg

tf_agents/bandits/agents/examples/v2/trainer.py

@@ -161,7 +161,7 @@ def baseline_reward_fn(observation, per_action_reward_fns):
     metrics += [tf_metrics.AverageReturnMultiMetric(
         reward_spec=environment.reward_spec(),
         batch_size=environment.batch_size)]
-  else:
+  if not isinstance(environment.reward_spec(), dict):
     metrics += [
         tf_metrics.AverageReturnMetric(batch_size=environment.batch_size)]
 

tf_agents/metrics/tf_metrics.py

@@ -177,8 +177,11 @@ def call(self, trajectory):
         tf.where(trajectory.is_first(), tf.zeros_like(self._return_accumulator),
                  self._return_accumulator))
 
-    # Update accumulator with received rewards.
-    self._return_accumulator.assign_add(trajectory.reward)
+    # Update accumulator with received rewards. We are summing over all
+    # non-batch dimensions in case the reward is a vector.
+    self._return_accumulator.assign_add(
+        tf.reduce_sum(
+            trajectory.reward, axis=range(1, len(trajectory.reward.shape))))
 
     # Add final returns to buffer.
     last_episode_indices = tf.squeeze(tf.where(trajectory.is_last()), axis=-1)

tf_agents/metrics/tf_metrics_test.py

@@ -328,6 +328,37 @@ def testChosenActionHistogram(self, run_mode):
       self.evaluate(metric.reset())
       self.assertEmpty(self.evaluate(metric.result()))
 
+  @parameterized.named_parameters([
+      ('testAverageReturnMetricVectorGraph', context.graph_mode, 6,
+       tensor_spec.TensorSpec((2,), tf.float32, 'r'), 18.0),
+      ('testAverageReturnMetricVectorEager', context.eager_mode, 6,
+       tensor_spec.TensorSpec((5,), tf.float32, 'r'), 45.0),])
+  def testAverageReturnMetricVector(self, run_mode, num_trajectories,
+                                    reward_spec, expected_result):
+    with run_mode():
+      trajectories = self._create_trajectories()
+      multi_trajectories = []
+      for traj in trajectories:
+        new_reward = tf.stack([traj.reward] * reward_spec.shape.as_list()[0],
+                              axis=1)
+        new_traj = trajectory.Trajectory(
+            step_type=traj.step_type,
+            observation=traj.observation,
+            action=traj.action,
+            policy_info=traj.policy_info,
+            next_step_type=traj.next_step_type,
+            reward=new_reward,
+            discount=traj.discount)
+        multi_trajectories.append(new_traj)
+
+      metric = tf_metrics.AverageReturnMetric(batch_size=2)
+      self.evaluate(tf.compat.v1.global_variables_initializer())
+      self.evaluate(metric.init_variables())
+      for i in range(num_trajectories):
+        self.evaluate(metric(multi_trajectories[i]))
+
+      self.assertAllClose(expected_result, self.evaluate(metric.result()))
+
   @parameterized.named_parameters([
       ('testAverageReturnMultiMetricGraph', context.graph_mode, 6,
        tensor_spec.TensorSpec((2,), tf.float32, 'r'), [9.0, 9.0]),