class_docstr: A RL task: environment and a collection of trajectories.
class_docstr_tok: ["A","RL","task",":","environment","and","a","collection","of","trajectories","."]
docstr: Configures a RL task.
  
  Args:
    env: Environment confirming to the gym.Env interface or a string,
      in which case `gym.make` will be called on this string to create an env.
    initial_trajectories: either a dict or list of Trajectories to use
      at start or an int, in which case that many trajectories are
      collected using a random policy to play in env. It can be also a string
      and then it should direct to the location where previously recorded
      trajectories are stored.
    gamma: float: discount factor for calculating returns.
    dm_suite: whether we are using the DeepMind suite or the gym interface
    random_starts: use random starts for training of Atari agents.
    max_steps: optional int: cut all trajectory slices at that many steps.
      The trajectory will be continued in the next epochs, up to `time_limit`.
    time_limit: optional int: stop all trajectories after that many steps (or
      after getting "done"). If `None`, use the same value as `max_steps`.
    timestep_to_np: a function that turns a timestep into a numpy array
      (ie., a tensor); if None, we just use the state of the timestep to
      represent it, but other representations (such as embeddings that include
      actions or serialized representations) can be passed here.
    num_stacked_frames: the number of stacked frames for Atari.
    n_replay_epochs: the size of the replay buffer expressed in epochs.
docstr_tok: ["Configures","a","RL","task",".","Args",":","env",":","Environment","confirming","to","the","gym.Env","interface","or","a","string",",","in","which","case","`","gym.make","`","will","be","called","on","this","string","to","create","an","env",".","initial_trajectories",":","either","a","dict","or","list","of","Trajectories","to","use","at","start","or","an","int",",","in","which","case","that","many","trajectories","are","collected","using","a","random","policy","to","play","in","env",".","It","can","be","also","a","string","and","then","it","should","direct","to","the","location","where","previously","recorded","trajectories","are","stored",".","gamma",":","float",":","discount","factor","for","calculating","returns",".","dm_suite",":","whether","we","are","using","the","DeepMind","suite","or","the","gym","interface","random_starts",":","use","random","starts","for","training","of","Atari","agents",".","max_steps",":","optional","int",":","cut","all","trajectory","slices","at","that","many","steps",".","The","trajectory","will","be","continued","in","the","next","epochs",",","up","to","`","time_limit","`",".","time_limit",":","optional","int",":","stop","all","trajectories","after","that","many","steps","(","or","after","getting","``","done","''",")",".","If","`","None","`",",","use","the","same","value","as","`","max_steps","`",".","timestep_to_np",":","a","function","that","turns","a","timestep","into","a","numpy","array","(","ie.",",","a","tensor",")",";","if","None",",","we","just","use","the","state","of","the","timestep","to","represent","it",",","but","other","representations","(","such","as","embeddings","that","include","actions","or","serialized","representations",")","can","be","passed","here",".","num_stacked_frames",":","the","number","of","stacked","frames","for","Atari",".","n_replay_epochs",":","the","size","of","the","replay","buffer","expressed","in","epochs","."]
code: def __init__(
      self,
      env=gin.REQUIRED,
      initial_trajectories=1,
      gamma=0.99,
      dm_suite=False,
      random_starts=True,
      max_steps=None,
      time_limit=None,
      timestep_to_np=None,
      num_stacked_frames=1,
      n_replay_epochs=1,
  ):
      r"""Configures a RL task.
  
      Args:
        env: Environment confirming to the gym.Env interface or a string,
          in which case `gym.make` will be called on this string to create an env.
        initial_trajectories: either a dict or list of Trajectories to use
          at start or an int, in which case that many trajectories are
          collected using a random policy to play in env. It can be also a string
          and then it should direct to the location where previously recorded
          trajectories are stored.
        gamma: float: discount factor for calculating returns.
        dm_suite: whether we are using the DeepMind suite or the gym interface
        random_starts: use random starts for training of Atari agents.
        max_steps: optional int: cut all trajectory slices at that many steps.
          The trajectory will be continued in the next epochs, up to `time_limit`.
        time_limit: optional int: stop all trajectories after that many steps (or
          after getting "done"). If `None`, use the same value as `max_steps`.
        timestep_to_np: a function that turns a timestep into a numpy array
          (ie., a tensor); if None, we just use the state of the timestep to
          represent it, but other representations (such as embeddings that include
          actions or serialized representations) can be passed here.
        num_stacked_frames: the number of stacked frames for Atari.
        n_replay_epochs: the size of the replay buffer expressed in epochs.
      """
      if isinstance(env, str):
          self._env_name = env
          if dm_suite:
              eval_env = None
              env = None
          else:
              env = gym.make(self._env_name)
              eval_env = gym.make(self._env_name)
      else:
          self._env_name = type(env).__name__
          eval_env = env
      self._env = env
      self._eval_env = eval_env
      self._dm_suite = dm_suite
      self._max_steps = max_steps
      if time_limit is None:
          time_limit = max_steps
      self._time_limit = time_limit
      self._gamma = gamma
      self._initial_trajectories = initial_trajectories
      self._last_observation = None
      self._n_steps_left = time_limit
      # Example trajectory for determining input/output shapes of the networks.
      self._example_trajectory = self.play(
          _random_policy(self.action_space), only_eval=True
      )
      # TODO(lukaszkaiser): find a better way to pass initial trajectories,
      # whether they are an explicit list, a file, or a number of random ones.
      if isinstance(initial_trajectories, int):
          initial_trajectories = [
              self.play(_random_policy(self.action_space))
              for _ in range(initial_trajectories)
          ]
      if isinstance(initial_trajectories, str):
          initial_trajectories = (
              self.load_initial_trajectories_from_path(
                  initial_trajectories_path=initial_trajectories
              )
          )
      if isinstance(initial_trajectories, list):
          if initial_trajectories:
              initial_trajectories = {0: initial_trajectories}
          else:
              initial_trajectories = {}
      self._timestep_to_np = timestep_to_np
      # Stored trajectories are indexed by epoch and within each epoch they
      # are stored in the order of generation so we can implement replay buffers.
      # TODO(lukaszkaiser): use dump_trajectories from BaseTrainer to allow
      # saving and reading trajectories from disk.
      self._trajectories = collections.defaultdict(list)
      self._trajectories.update(initial_trajectories)
      # When we repeatedly save, trajectories for many epochs do not change, so
      # we don't need to save them again. This keeps track which are unchanged.
      self._saved_epochs_unchanged = []
      self._n_replay_epochs = n_replay_epochs
      self._n_trajectories = 0
      self._n_interactions = 0
  
imports: import collections
  import os
  import gin
  import gym
  import numpy
  from trax import fastmath
  from trax.rl import advantages
  from trax.supervised import training
label_desc: Use image node_id 1 to create a new RLTask object  with example: obj = RLTask(env, initial_trajectories, gamma, dm_suite, random_starts, max_steps, time_limit, timestep_to_np, num_stacked_frames, n_replay_epochs)