code: def _add_next_action_extras(
      double_transitions: Transition,
  ) -> reverb.ReplaySample:
      # As TD3 is online by default, it expects an iterator over replay samples.
      info = tree.map_structure(
          lambda dtype: tf.ones([], dtype),
          reverb.SampleInfo.tf_dtypes(),
      )
      return reverb.ReplaySample(
          info=info,
          data=Transition(
              observation=double_transitions.observation[0],
              action=double_transitions.action[0],
              reward=double_transitions.reward[0],
              discount=double_transitions.discount[0],
              next_observation=double_transitions.next_observation[0],
              extras={"next_action": double_transitions.action[1]},
          ),
      )
  
imports: from absl import app
  from absl import flags
  import acme
  from acme import specs
  from acme.agents.jax import actor_core
  from acme.agents.jax import actors
  from acme.agents.jax import td3
  from acme.datasets import tfds
  from acme.examples.offline import helpers
  from acme.jax import variable_utils
  from acme.types import Transition
  from acme.utils import loggers
  import haiku
  import jax
  import optax
  import reverb
  import rlds
  import tensorflow
  import tree
label_desc: Use image node_id 1 for calling a global function with example usage: _add_next_action_extras(double_transitions) and returns: reverb