code: def wrap_experiment(
      function=None,
      *,
      log_dir=None,
      prefix="experiment",
      name=None,
      snapshot_mode="last",
      snapshot_gap=1,
      archive_launch_repo=True,
      name_parameters=None,
      use_existing_dir=False,
      x_axis="TotalEnvSteps",
  ):
      """Decorate a function to turn it into an ExperimentTemplate.
  
      When invoked, the wrapped function will receive an ExperimentContext, which
      will contain the log directory into which the experiment should log
      information.
  
      This decorator can be invoked in two differed ways.
  
      Without arguments, like this:
  
          @wrap_experiment
          def my_experiment(ctxt, seed, lr=0.5):
              ...
  
      Or with arguments:
  
          @wrap_experiment(snapshot_mode='all')
          def my_experiment(ctxt, seed, lr=0.5):
              ...
  
      All arguments must be keyword arguments.
  
      Args:
          function (callable or None): The experiment function to wrap.
          log_dir (str or None): The full log directory to log to. Will be
              computed from `name` if omitted.
          name (str or None): The name of this experiment template. Will be
              filled from the wrapped function's name if omitted.
          prefix (str): Directory under data/local in which to place the
              experiment directory.
          snapshot_mode (str): Policy for which snapshots to keep (or make at
              all). Can be either "all" (all iterations will be saved), "last"
              (only the last iteration will be saved), "gap" (every snapshot_gap
              iterations are saved), or "none" (do not save snapshots).
          snapshot_gap (int): Gap between snapshot iterations. Waits this number
              of iterations before taking another snapshot.
          archive_launch_repo (bool): Whether to save an archive of the
              repository containing the launcher script. This is a potentially
              expensive operation which is useful for ensuring reproducibility.
          name_parameters (str or None): Parameters to insert into the experiment
              name. Should be either None (the default), 'all' (all parameters
              will be used), or 'passed' (only passed parameters will be used).
              The used parameters will be inserted in the order they appear in
              the function definition.
          use_existing_dir (bool): If true, (re)use the directory for this
              experiment, even if it already contains data.
          x_axis (str): Key to use for x axis of plots.
  
      Returns:
          callable: The wrapped function.
  
      """
      return ExperimentTemplate(
          function=function,
          log_dir=log_dir,
          prefix=prefix,
          name=name,
          snapshot_mode=snapshot_mode,
          snapshot_gap=snapshot_gap,
          archive_launch_repo=archive_launch_repo,
          name_parameters=name_parameters,
          use_existing_dir=use_existing_dir,
          x_axis=x_axis,
      )
  
imports: import collections
  import datetime
  import enum
  import functools
  import gc
  import inspect
  import json
  import os
  import pathlib
  import subprocess
  import warnings
  import weakref
  import dateutil.tz
  import dowel
  from dowel import logger
  import numpy
  import __main__
label_desc: Use image node_id 3 for calling a global function with example usage: wrap_experiment(function) and returns: ExperimentTemplate