text: import numpy as np
  
  import rllab.misc.logger as logger
  import rllab.plotter as plotter
  from rllab.algos.base import RLAlgorithm
  from rllab.policies.base import Policy
  from rllab.sample_processors.default_sample_processor import DefaultSampleProcessor
  #from rllab.sample_processors.vectorized_sampler import VectorizedSampler
  from rllab.sampler import parallel_sampler
  from rllab.sampler.base import BaseSampler
  
  
  class BatchSampler(BaseSampler):
      def __init__(self, algo):
          """
          :type algo: BatchPolopt
          """
          self.algo = algo
  
      def start_worker(self):
          parallel_sampler.populate_task(self.algo.env, self.algo.policy, scope=self.algo.scope)
  
      def shutdown_worker(self):
          parallel_sampler.terminate_task(scope=self.algo.scope)
  
      def obtain_samples(self, itr):
          cur_params = self.algo.policy.get_param_values()
          paths = parallel_sampler.sample_paths(
              policy_params=cur_params,
              max_samples=self.algo.batch_size,
              max_path_length=self.algo.max_path_length,
              scope=self.algo.scope,
          )
          if self.algo.whole_paths:
              return paths
          else:
              paths_truncated = parallel_sampler.truncate_paths(paths, self.algo.batch_size)
              return paths_truncated
  
  
  class BatchPolopt(RLAlgorithm):
      """
      Base class for batch sampling-based policy optimization methods.
      This includes various policy gradient methods like vpg, npg, ppo, trpo, etc.
      """
  
      def __init__(
              self,
              env,
              policy,
              baseline,
              scope=None,
              n_itr=500,
              start_itr=0,
              batch_size=5000,
              max_path_length=500,
              discount=0.99,
              gae_lambda=1,
              plot=False,
              pause_for_plot=False,
              center_adv=True,
              positive_adv=False,
              store_paths=False,
              whole_paths=True,
              sampler_cls=None,
              sampler_args=None,
              sample_processor_cls=None,
              sample_processor_args=None,
              n_vectorized_envs=None,
              parallel_vec_env=False,
              **kwargs
      ):
          """
          :param env: Environment
          :param policy: Policy
          :type policy: Policy
          :param baseline: Baseline
          :param scope: Scope for identifying the algorithm. Must be specified if running multiple algorithms
          simultaneously, each using different environments and policies
          :param n_itr: Number of iterations.
          :param start_itr: Starting iteration.
          :param batch_size: Number of samples per iteration.
          :param max_path_length: Maximum length of a single rollout.
          :param discount: Discount.
          :param gae_lambda: Lambda used for generalized advantage estimation.
          :param plot: Plot evaluation run after each iteration.
          :param pause_for_plot: Whether to pause before contiuing when plotting.
          :param center_adv: Whether to rescale the advantages so that they have mean 0 and standard deviation 1.
          :param positive_adv: Whether to shift the advantages so that they are always positive. When used in
          conjunction with center_adv the advantages will be standardized before shifting.
          :param store_paths: Whether to save all paths data to the snapshot.
          """
          self.env = env
          self.policy = policy
          self.baseline = baseline
          self.scope = scope
          self.n_itr = n_itr
          self.current_itr = start_itr
          self.batch_size = batch_size
          self.max_path_length = max_path_length
          self.discount = discount
          self.gae_lambda = gae_lambda
          self.plot = plot
          self.pause_for_plot = pause_for_plot
          self.center_adv = center_adv
          self.positive_adv = positive_adv
          self.store_paths = store_paths
          self.whole_paths = whole_paths
          if sampler_cls is None:
              # if n_vectorized_envs is None:
              #     n_vectorized_envs = min(100, max(1, int(np.ceil(batch_size / max_path_length))))
  
              # FIXME commenting out the line below might break things
              # if self.policy.vectorized:
              #     self.sampler = VectorizedSampler(env=env, policy=policy, n_envs=n_vectorized_envs, parallel=parallel_vec_env)
              # else:
              if sampler_args is None:
                  sampler_args = dict()
              self.sampler = BatchSampler(self, **sampler_args)
  
          if sample_processor_cls is None:
              sample_processor_cls = DefaultSampleProcessor
          if sample_processor_args is None:
              sample_processor_args = dict(algo=self)
          else:
              sample_processor_args = dict(sample_processor_args, algo=self)
  
          self.sample_processor = sample_processor_cls(**sample_processor_args)
  
      def start_worker(self):
          self.sampler.start_worker()
          if self.plot:
              plotter.init_plot(self.env, self.policy)
  
      def shutdown_worker(self):
          self.sampler.shutdown_worker()
  
      def obtain_samples(self, itr):
          # if isinstance(self.sampler, VectorizedSampler):
          #     return self.sampler.obtain_samples(
          #         max_path_length=self.max_path_length,
          #         batch_size=self.batch_size
          #     )
          # else:
              return self.sampler.obtain_samples(itr)
  
      def process_samples(self, itr, paths):
          if self.policy.vectorized:
              return self.sample_processor.process_samples(itr, paths)
          else:
              return self.sampler.process_samples(itr, paths)
  
      def train(self):
          self.start_worker()
          self.init_opt()
          for itr in range(self.current_itr, self.n_itr):
              with logger.prefix('itr #%d | ' % itr):
                  paths = self.obtain_samples(itr)
                  samples_data = self.process_samples(itr, paths)
                  # TOFIX(eugene) why is this here, and can I get rid of it?
                  self.log_diagnostics(paths)
                  self.optimize_policy(itr, samples_data)
                  logger.log("saving snapshot...")
                  params = self.get_itr_snapshot(itr, samples_data)
                  self.current_itr = itr + 1
                  # FIXME(eugene) uncomment this line
                  #params["algo"] = self
                  if self.store_paths:
                      params["paths"] = samples_data["paths"]
                  logger.save_itr_params(itr, params)
                  logger.log("saved")
                  logger.dump_tabular(with_prefix=False)
                  if self.plot:
                      self.update_plot()
                      if self.pause_for_plot:
                          input("Plotting evaluation run: Press Enter to "
                                    "continue...")
  
          self.shutdown_worker()
  
      def log_diagnostics(self, paths):
          self.env.log_diagnostics(paths)
          self.policy.log_diagnostics(paths)
          # FIXME(eugene) support this later
          #self.baseline.log_diagnostics(paths)
  
      def init_opt(self):
          """
          Initialize the optimization procedure. If using theano / cgt, this may
          include declaring all the variables and compiling functions
          """
          raise NotImplementedError
  
      def get_itr_snapshot(self, itr, samples_data):
          """
          Returns all the data that should be saved in the snapshot for this
          iteration.
          """
          raise NotImplementedError
  
      def optimize_policy(self, itr, samples_data):
          raise NotImplementedError
  
      def update_plot(self):
          if self.plot:
              plotter.update_plot(self.policy, self.max_path_length)
  