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).
docstr: Perform algorithm optimizing.
  
  Args:
      indices (list): Tasks used for training.
code: def _optimize_policy(self, indices):
      """Perform algorithm optimizing.
  
      Args:
          indices (list): Tasks used for training.
  
      """
      num_tasks = len(indices)
      context = self._sample_context(indices)
      # clear context and reset belief of policy
      self._policy.reset_belief(num_tasks=num_tasks)
  
      # data shape is (task, batch, feat)
      obs, actions, rewards, next_obs, terms = self._sample_data(
          indices
      )
      policy_outputs, task_z = self._policy(obs, context)
      new_actions, policy_mean, policy_log_std, log_pi = policy_outputs[
          :4
      ]
  
      # flatten out the task dimension
      t, b, _ = obs.size()
      obs = obs.view(t * b, -1)
      actions = actions.view(t * b, -1)
      next_obs = next_obs.view(t * b, -1)
  
      # optimize qf and encoder networks
      q1_pred = self._qf1(torch.cat([obs, actions], dim=1), task_z)
      q2_pred = self._qf2(torch.cat([obs, actions], dim=1), task_z)
      v_pred = self._vf(obs, task_z.detach())
  
      with torch.no_grad():
          target_v_values = self.target_vf(next_obs, task_z)
  
      # KL constraint on z if probabilistic
      zero_optim_grads(self.context_optimizer)
      if self._use_information_bottleneck:
          kl_div = self._policy.compute_kl_div()
          kl_loss = self._kl_lambda * kl_div
          kl_loss.backward(retain_graph=True)
  
      zero_optim_grads(self.qf1_optimizer)
      zero_optim_grads(self.qf2_optimizer)
  
      rewards_flat = rewards.view(self._batch_size * num_tasks, -1)
      rewards_flat = rewards_flat * self._reward_scale
      terms_flat = terms.view(self._batch_size * num_tasks, -1)
      q_target = (
          rewards_flat
          + (1.0 - terms_flat) * self._discount * target_v_values
      )
      qf_loss = torch.mean((q1_pred - q_target) ** 2) + torch.mean(
          (q2_pred - q_target) ** 2
      )
      qf_loss.backward()
  
      self.qf1_optimizer.step()
      self.qf2_optimizer.step()
      self.context_optimizer.step()
  
      # compute min Q on the new actions
      q1 = self._qf1(
          torch.cat([obs, new_actions], dim=1), task_z.detach()
      )
      q2 = self._qf2(
          torch.cat([obs, new_actions], dim=1), task_z.detach()
      )
      min_q = torch.min(q1, q2)
  
      # optimize vf
      v_target = min_q - log_pi
      vf_loss = self.vf_criterion(v_pred, v_target.detach())
      zero_optim_grads(self.vf_optimizer)
      vf_loss.backward()
      self.vf_optimizer.step()
      self._update_target_network()
  
      # optimize policy
      log_policy_target = min_q
      policy_loss = (log_pi - log_policy_target).mean()
  
      mean_reg_loss = (
          self._policy_mean_reg_coeff * (policy_mean**2).mean()
      )
      std_reg_loss = (
          self._policy_std_reg_coeff * (policy_log_std**2).mean()
      )
      pre_tanh_value = policy_outputs[-1]
      pre_activation_reg_loss = self._policy_pre_activation_coeff * (
          (pre_tanh_value**2).sum(dim=1).mean()
      )
      policy_reg_loss = (
          mean_reg_loss + std_reg_loss + pre_activation_reg_loss
      )
      policy_loss = policy_loss + policy_reg_loss
  
      zero_optim_grads(self._policy_optimizer)
      policy_loss.backward()
      self._policy_optimizer.step()
  
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 6 for calling the PEARL obj's underlying member method code with example usage: obj._optimize_policy(indices) without return types