class_docstr: A PEARL model based on https://arxiv.org/abs/1903.08254.
  
  PEARL, which stands for Probablistic Embeddings for Actor-Critic
  Reinforcement Learning, is an off-policy meta-RL algorithm. It is built
  on top of SAC using two Q-functions and a value function with an addition
  of an inference network that estimates the posterior :math:`q(z \| c)`.
  The policy is conditioned on the latent variable Z in order to adpat its
  behavior to specific tasks.
  
  Args:
      env (list[Environment]): Batch of sampled environment updates(
          EnvUpdate), which, when invoked on environments, will configure
          them with new tasks.
      policy_class (type): Class implementing
          :pyclass:`~ContextConditionedPolicy`
      encoder_class (garage.torch.embeddings.ContextEncoder): Encoder class
          for the encoder in context-conditioned policy.
      inner_policy (garage.torch.policies.Policy): Policy.
      qf (torch.nn.Module): Q-function.
      vf (torch.nn.Module): Value function.
      sampler (garage.sampler.Sampler): Sampler.
      num_train_tasks (int): Number of tasks for training.
      num_test_tasks (int or None): Number of tasks for testing.
      latent_dim (int): Size of latent context vector.
      encoder_hidden_sizes (list[int]): Output dimension of dense layer(s) of
          the context encoder.
      test_env_sampler (garage.experiment.SetTaskSampler): Sampler for test
          tasks.
      policy_lr (float): Policy learning rate.
      qf_lr (float): Q-function learning rate.
      vf_lr (float): Value function learning rate.
      context_lr (float): Inference network learning rate.
      policy_mean_reg_coeff (float): Policy mean regulation weight.
      policy_std_reg_coeff (float): Policy std regulation weight.
      policy_pre_activation_coeff (float): Policy pre-activation weight.
      soft_target_tau (float): Interpolation parameter for doing the
          soft target update.
      kl_lambda (float): KL lambda value.
      optimizer_class (type): Type of optimizer for training networks.
      use_information_bottleneck (bool): False means latent context is
          deterministic.
      use_next_obs_in_context (bool): Whether or not to use next observation
          in distinguishing between tasks.
      meta_batch_size (int): Meta batch size.
      num_steps_per_epoch (int): Number of iterations per epoch.
      num_initial_steps (int): Number of transitions obtained per task before
          training.
      num_tasks_sample (int): Number of random tasks to obtain data for each
          iteration.
      num_steps_prior (int): Number of transitions to obtain per task with
          z ~ prior.
      num_steps_posterior (int): Number of transitions to obtain per task
          with z ~ posterior.
      num_extra_rl_steps_posterior (int): Number of additional transitions
          to obtain per task with z ~ posterior that are only used to train
          the policy and NOT the encoder.
      batch_size (int): Number of transitions in RL batch.
      embedding_batch_size (int): Number of transitions in context batch.
      embedding_mini_batch_size (int): Number of transitions in mini context
          batch; should be same as embedding_batch_size for non-recurrent
          encoder.
      discount (float): RL discount factor.
      replay_buffer_size (int): Maximum samples in replay buffer.
      reward_scale (int): Reward scale.
      update_post_train (int): How often to resample context when obtaining
          data during training (in episodes).
code: def __init__(
      self,
      env,
      inner_policy,
      qf,
      vf,
      sampler,
      *,  # Mostly numbers after here.
      num_train_tasks,
      num_test_tasks=None,
      latent_dim,
      encoder_hidden_sizes,
      test_env_sampler,
      policy_class=ContextConditionedPolicy,
      encoder_class=MLPEncoder,
      policy_lr=3e-4,
      qf_lr=3e-4,
      vf_lr=3e-4,
      context_lr=3e-4,
      policy_mean_reg_coeff=1e-3,
      policy_std_reg_coeff=1e-3,
      policy_pre_activation_coeff=0.0,
      soft_target_tau=0.005,
      kl_lambda=0.1,
      optimizer_class=torch.optim.Adam,
      use_information_bottleneck=True,
      use_next_obs_in_context=False,
      meta_batch_size=64,
      num_steps_per_epoch=1000,
      num_initial_steps=100,
      num_tasks_sample=100,
      num_steps_prior=100,
      num_steps_posterior=0,
      num_extra_rl_steps_posterior=100,
      batch_size=1024,
      embedding_batch_size=1024,
      embedding_mini_batch_size=1024,
      discount=0.99,
      replay_buffer_size=1000000,
      reward_scale=1,
      update_post_train=1,
  ):
      self._env = env
      self._qf1 = qf
      self._qf2 = copy.deepcopy(qf)
      self._vf = vf
      self._num_train_tasks = num_train_tasks
      self._latent_dim = latent_dim
  
      self._policy_mean_reg_coeff = policy_mean_reg_coeff
      self._policy_std_reg_coeff = policy_std_reg_coeff
      self._policy_pre_activation_coeff = policy_pre_activation_coeff
      self._soft_target_tau = soft_target_tau
      self._kl_lambda = kl_lambda
      self._use_information_bottleneck = use_information_bottleneck
      self._use_next_obs_in_context = use_next_obs_in_context
  
      self._meta_batch_size = meta_batch_size
      self._num_steps_per_epoch = num_steps_per_epoch
      self._num_initial_steps = num_initial_steps
      self._num_tasks_sample = num_tasks_sample
      self._num_steps_prior = num_steps_prior
      self._num_steps_posterior = num_steps_posterior
      self._num_extra_rl_steps_posterior = num_extra_rl_steps_posterior
      self._batch_size = batch_size
      self._embedding_batch_size = embedding_batch_size
      self._embedding_mini_batch_size = embedding_mini_batch_size
      self._discount = discount
      self._replay_buffer_size = replay_buffer_size
      self._reward_scale = reward_scale
      self._update_post_train = update_post_train
      self._task_idx = None
      self._single_env = env[0]()
      self.max_episode_length = self._single_env.spec.max_episode_length
  
      self._sampler = sampler
  
      self._is_resuming = False
  
      if num_test_tasks is None:
          num_test_tasks = test_env_sampler.n_tasks
      if num_test_tasks is None:
          raise ValueError(
              "num_test_tasks must be provided if "
              "test_env_sampler.n_tasks is None"
          )
  
      worker_args = dict(deterministic=True, accum_context=True)
      self._evaluator = MetaEvaluator(
          test_task_sampler=test_env_sampler,
          worker_class=PEARLWorker,
          worker_args=worker_args,
          n_test_tasks=num_test_tasks,
      )
  
      encoder_spec = self.get_env_spec(
          self._single_env, latent_dim, "encoder"
      )
      encoder_in_dim = int(np.prod(encoder_spec.input_space.shape))
      encoder_out_dim = int(np.prod(encoder_spec.output_space.shape))
      context_encoder = encoder_class(
          input_dim=encoder_in_dim,
          output_dim=encoder_out_dim,
          hidden_sizes=encoder_hidden_sizes,
      )
  
      self._policy = policy_class(
          latent_dim=latent_dim,
          context_encoder=context_encoder,
          policy=inner_policy,
          use_information_bottleneck=use_information_bottleneck,
          use_next_obs=use_next_obs_in_context,
      )
  
      # buffer for training RL update
      self._replay_buffers = {
          i: PathBuffer(replay_buffer_size)
          for i in range(num_train_tasks)
      }
  
      self._context_replay_buffers = {
          i: PathBuffer(replay_buffer_size)
          for i in range(num_train_tasks)
      }
  
      self.target_vf = copy.deepcopy(self._vf)
      self.vf_criterion = torch.nn.MSELoss()
  
      self._policy_optimizer = optimizer_class(
          self._policy.networks[1].parameters(),
          lr=policy_lr,
      )
      self.qf1_optimizer = optimizer_class(
          self._qf1.parameters(),
          lr=qf_lr,
      )
      self.qf2_optimizer = optimizer_class(
          self._qf2.parameters(),
          lr=qf_lr,
      )
      self.vf_optimizer = optimizer_class(
          self._vf.parameters(),
          lr=vf_lr,
      )
      self.context_optimizer = optimizer_class(
          self._policy.networks[0].parameters(),
          lr=context_lr,
      )
  
imports: import copy
  import akro
  from dowel import logger
  import numpy
  import torch
  from garage import EnvSpec, InOutSpec, StepType, TimeStep
  from garage.experiment import MetaEvaluator
  from garage.np.algos import MetaRLAlgorithm
  from garage.replay_buffer import PathBuffer
  from garage.sampler import DefaultWorker
  from garage.torch import global_device
  from garage.torch._functions import np_to_torch, zero_optim_grads
  from garage.torch.embeddings import MLPEncoder
  from garage.torch.policies import ContextConditionedPolicy
label_desc: Use image node_id 1 to create a new PEARL object from inherited base classes: MetaRLAlgorithm with example: obj = PEARL(env, inner_policy, qf, vf, sampler)