Incorrect formatting for time_step, trajectory when using tf.print

[coreyleveen]

When passing a TimeStep or Trajectory to tf.print, the output presents the objects incorrectly in that some field values get printed in the wrong places.

The following code demonstrates this:

import tensorflow as tf

from tf_agents.trajectories import time_step
from tf_agents.trajectories import trajectory

ts = time_step.TimeStep(
  tf.constant(0),     # step_type
  tf.constant(0.5),   # reward
  tf.constant(0.99),  # discount
  tf.constant([1, 1]) # observation
)

tf.print(ts)
#=> TimeStep(
#{'discount': 0, 'observation': 0.5, 'reward': 0.99, 'step_type': [1 1]})

observation, action, policy_info = (), (), ()
reward = tf.constant([2.0, 2.0, 2.0])
discount = tf.constant([0.99, 0.99, 0.99])
traj = trajectory.first(observation, action, policy_info, reward, discount)

tf.print(traj)
#=> Trajectory(
#{'action': (),
# 'discount': [0 0 0],
# 'next_step_type': [1 1 1],
# 'observation': (),
# 'policy_info': (),
# 'reward': [2 2 2],
# 'step_type': [0.99 0.99 0.99]})

Note the incorrect values in the output of the above code.
This occurs because both TimeStep and Trajectory use pprint.pformat in their respective __repr__ method, which sorts by the dict keys when printing. However, tf.print still prints field values in the order of the namedtuple's _fields, which results in a mismatch of field names and the associated values.