code: def _define_train(
      train_env,
      ppo_hparams,
      eval_env_fn=None,
      sampling_temp=1.0,
      distributional_size=1,
      distributional_subscale=0.04,
      distributional_threshold=0.0,
      epoch=-1,
      **collect_kwargs,
  ):
      """Define the training setup."""
      memory, collect_summary, train_initialization = _define_collect(
          train_env,
          ppo_hparams,
          "ppo_train",
          eval_phase=False,
          sampling_temp=sampling_temp,
          distributional_size=distributional_size,
          **collect_kwargs,
      )
      ppo_summary = ppo.define_ppo_epoch(
          memory,
          ppo_hparams,
          train_env.action_space,
          train_env.batch_size,
          distributional_size=distributional_size,
          distributional_subscale=distributional_subscale,
          distributional_threshold=distributional_threshold,
          epoch=epoch,
      )
      train_summary = tf.summary.merge([collect_summary, ppo_summary])
  
      if ppo_hparams.eval_every_epochs:
          # TODO(koz4k): Do we need this at all?
          assert eval_env_fn is not None
          eval_env = eval_env_fn(in_graph=True)
          (
              _,
              eval_collect_summary,
              eval_initialization,
          ) = _define_collect(
              eval_env,
              ppo_hparams,
              "ppo_eval",
              eval_phase=True,
              sampling_temp=0.0,
              distributional_size=distributional_size,
              **collect_kwargs,
          )
          return (
              train_summary,
              eval_collect_summary,
              (train_initialization, eval_initialization),
          )
      else:
          return (train_summary, None, (train_initialization,))
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import math
  import os
  from tensor2tensor.layers import common_layers
  from tensor2tensor.models.research.rl import get_policy
  from tensor2tensor.rl import ppo
  from tensor2tensor.rl.envs.tf_atari_wrappers import StackWrapper
  from tensor2tensor.rl.envs.tf_atari_wrappers import WrapperBase
  from tensor2tensor.rl.policy_learner import PolicyLearner
  from tensor2tensor.rl.restarter import Restarter
  from tensor2tensor.utils import trainer_lib
  import tensorflow.compat.v1
  import tensorflow_probability
label_desc: Use image node_id 1 for calling a global function with example usage: _define_train(train_env, ppo_hparams, eval_env_fn, sampling_temp, distributional_size, distributional_subscale, distributional_threshold, epoch) and returns: train_summary,  eval_collect_summary,  train_initialization,  eval_initialization, train_summary,  None,  train_initialization