code: def from_episode(
      observation: types.NestedSpecTensorOrArray,
      action: types.NestedSpecTensorOrArray,
      policy_info: types.NestedSpecTensorOrArray,
      reward: types.NestedSpecTensorOrArray,
      discount: Optional[types.SpecTensorOrArray] = None,
  ) -> Trajectory:
      """Create a Trajectory from tensors representing a single episode.
  
      If none of the inputs are tensors, then numpy arrays are generated instead.
  
      If `discount` is not provided, the first entry in `reward` is used to estimate
      `T`:
  
      ```
      reward_0 = tf.nest.flatten(reward)[0]
      T = shape(reward_0)[0]
      ```
  
      In this case, a `discount` of all ones having dtype `float32` is generated.
  
      **NOTE**: all tensors/numpy arrays passed to this function have the same time
      dimension `T`. When the generated trajectory passes through `to_transition`,
      it will only return a `(time_steps, next_time_steps)` pair with `T - 1` in the
      time dimension, which means the reward at step T is dropped. So if the reward
      at step `T` is important, please make sure the episode passed to this function
      contains an additional step.
  
      Args:
        observation: (possibly nested tuple of) `Tensor` or `np.ndarray`; all shaped
          `[T, ...]`.
        action: (possibly nested tuple of) `Tensor` or `np.ndarray`; all shaped `[T,
          ...]`.
        policy_info: (possibly nested tuple of) `Tensor` or `np.ndarray`; all shaped
          `[T, ...]`.
        reward: (possibly nested tuple of) `Tensor` or `np.ndarray`; all shaped `[T,
          ...]`.
        discount: A floating point vector `Tensor` or `np.ndarray`; shaped `[T]`
          (optional).
  
      Returns:
        An instance of `Trajectory`.
      """
      use_tensors = nest_utils.has_tensors(
          observation, action, policy_info, reward, discount
      )
      map_structure = functools.partial(
          tf.nest.map_structure, expand_composites=True
      )
      if use_tensors:
          ones_fn = tf.ones
          float_dtype = tf.float32
          convert_fn = tf.convert_to_tensor
          concat_fn = tf.concat
          maximum_fn = tf.maximum
          fill_fn = tf.fill
          identity_map = lambda struct: map_structure(
              tf.identity, struct
          )
      else:
          ones_fn = np.ones
          float_dtype = np.float32
          convert_fn = np.asarray
          concat_fn = np.concatenate
          maximum_fn = np.maximum
          fill_fn = np.full
          identity_map = lambda struct: map_structure(
              np.asarray, struct
          )
  
      def _from_episode(
          observation, action, policy_info, reward, discount
      ):
          """Implementation of from_episode."""
          if discount is not None:
              time_source = discount
          else:
              time_source = tf.nest.flatten(reward)[0]
          if tf.is_tensor(time_source):
              num_frames = _maybe_static_outer_dim(time_source)
          else:
              num_frames = np.shape(time_source)[0]
          if discount is None:
              discount = ones_fn([num_frames], dtype=float_dtype)
  
          if not tf.is_tensor(num_frames):
  
              def check_num_frames(t):
                  if tf.is_tensor(t):
                      outer_dim = _maybe_static_outer_dim(t)
                  else:
                      outer_dim = t.shape[0]
                  if (
                      not tf.is_tensor(outer_dim)
                      and outer_dim != num_frames
                  ):
                      raise ValueError(
                          "Expected first dimension to be {}, but saw outer dim: {}".format(
                              num_frames, outer_dim
                          )
                      )
  
              tf.nest.map_structure(
                  check_num_frames,
                  (observation, action, policy_info, reward, discount),
                  expand_composites=False,
              )
  
          ts_first = convert_fn(ts.StepType.FIRST)
          ts_last = convert_fn(ts.StepType.LAST)
          mid_size = maximum_fn(0, num_frames - 1)
          ts_mid = fill_fn([mid_size], ts.StepType.MID)
          step_type = concat_fn(([ts_first], ts_mid), axis=0)
          next_step_type = concat_fn((ts_mid, [ts_last]), axis=0)
  
          return Trajectory(
              step_type=step_type,
              observation=identity_map(observation),
              action=identity_map(action),
              policy_info=identity_map(policy_info),
              next_step_type=next_step_type,
              reward=identity_map(reward),
              discount=identity_map(discount),
          )
  
      if use_tensors:
          with tf.name_scope("from_episode"):
              return _from_episode(
                  observation, action, policy_info, reward, discount
              )
      else:
          return _from_episode(
              observation, action, policy_info, reward, discount
          )
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import functools
  import pprint
  from typing import NamedTuple, Optional
  import numpy
  import tensorflow
  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 composite
  from tf_agents.utils import nest_utils
  from tf_agents.utils import value_ops
label_desc: Use image node_id 8 for calling a global function with example usage: from_episode(observation, action, policy_info, reward, discount) and returns: Trajectory, _from_episode, _from_episode