code: def _convert_timestep(timestep: ts.TimeStep) -> dm_env.TimeStep:
      if timestep.step_type.is_first():
          step_type = dm_env.StepType.FIRST
      elif timestep.step_type.is_last():
          step_type = dm_env.StepType.LAST
      elif timestep.step_type.is_mid():
          step_type = dm_env.StepType.MID
      else:
          raise ValueError(f"Invalid Step type: {timestep.step_type}")
      return dm_env.TimeStep(
          step_type=step_type,
          reward=None if step_type.first() else timestep.reward,
          discount=None if step_type.first() else timestep.discount,
          observation=timestep.observation,
      )
  
file_path: agents/tf_agents/environments/py_to_dm_wrapper.py
imports: from typing import Text
  import dm_env
  from dm_env import specs
  from tf_agents.environments import py_environment
  from tf_agents.specs import array_spec
  from tf_agents.trajectories import time_step
  import tree
label_desc: Use image node_id 1 for calling a global function with example usage: _convert_timestep(timestep) and returns: dm_env