code: def preprocess_step(timestep: dm_env.TimeStep) -> dm_env.TimeStep:
      if timestep.discount is None:
          timestep = timestep._replace(discount=1.0)
      return tree.map_structure(_preprocess_none, timestep)
  
file_path: dm-haiku/examples/impala/util.py
imports: import collections
  from absl import logging
  import dm_env
  import numpy
  import tree
label_desc: Use image node_id 2 for calling a global function with example usage: preprocess_step(timestep) and returns: tree