docstr: Args:
      obs: An float32 tensor of shape ([B] + observation_space).
           E.g. [B, C, H, W] in atari.
      actions: An int64 tensor of shape [B].
      behaviour_logits: A float32 tensor of shape [B, NUM_ACTIONS].
      rewards: A float32 tensor of shape [B].
      dones: A float32 tensor of shape [B].
      learning_rate: float scalar of learning rate.
      entropy_coeff: float scalar of entropy coefficient.
docstr_tok: ["Args",":","obs",":","An","float32","tensor","of","shape","(","[","B","]","+","observation_space",")",".","E.g",".","[","B",",","C",",","H",",","W","]","in","atari",".","actions",":","An","int64","tensor","of","shape","[","B","]",".","behaviour_logits",":","A","float32","tensor","of","shape","[","B",",","NUM_ACTIONS","]",".","rewards",":","A","float32","tensor","of","shape","[","B","]",".","dones",":","A","float32","tensor","of","shape","[","B","]",".","learning_rate",":","float","scalar","of","learning","rate",".","entropy_coeff",":","float","scalar","of","entropy","coefficient","."]
code: def learn(
      self,
      obs,
      actions,
      behaviour_logits,
      rewards,
      dones,
      learning_rate,
      entropy_coeff,
  ):
      """
      Args:
          obs: An float32 tensor of shape ([B] + observation_space).
               E.g. [B, C, H, W] in atari.
          actions: An int64 tensor of shape [B].
          behaviour_logits: A float32 tensor of shape [B, NUM_ACTIONS].
          rewards: A float32 tensor of shape [B].
          dones: A float32 tensor of shape [B].
          learning_rate: float scalar of learning rate.
          entropy_coeff: float scalar of entropy coefficient.
      """
  
      values = self.model.value(obs)
      target_logits = self.model.policy(obs)
  
      target_policy_distribution = Categorical(target_logits)
      behaviour_policy_distribution = Categorical(behaviour_logits)
  
      policy_entropy = target_policy_distribution.entropy()
  
      target_actions_log_probs = self._log_prob(target_logits, actions)
      behaviour_actions_log_probs = self._log_prob(
          behaviour_logits, actions
      )
  
      # Calculating kl for debug
      kl = target_policy_distribution.kl_divergence(
          behaviour_policy_distribution
      )
      kl = paddle.mean(kl)
      """
          Split the tensor into batches at known episode cut boundaries. 
          [B * T] -> [T, B]
          """
      T = self.sample_batch_steps
      B = obs.shape[0] // T
  
      def split_batches(tensor):
          splited_tensor = paddle.reshape(
              tensor, [B, T] + list(tensor.shape[1:])
          )
          # transpose B and T
          return paddle.transpose(
              splited_tensor,
              [1, 0] + list(range(2, 1 + len(tensor.shape))),
          )
  
      behaviour_actions_log_probs = split_batches(
          behaviour_actions_log_probs
      )
      target_actions_log_probs = split_batches(target_actions_log_probs)
      policy_entropy = split_batches(policy_entropy)
      dones = split_batches(dones)
      rewards = split_batches(rewards)
      values = split_batches(values)
  
      # [T, B] -> [T - 1, B] for V-trace calc.
      behaviour_actions_log_probs = behaviour_actions_log_probs[:-1]
      target_actions_log_probs = target_actions_log_probs[:-1]
      policy_entropy = policy_entropy[:-1]
      dones = dones[:-1]
      rewards = rewards[:-1]
      bootstrap_value = values[-1:]
      values = values[:-1]
  
      bootstrap_value = paddle.squeeze(bootstrap_value, 0)
  
      vtrace_loss = VTraceLoss(
          behaviour_actions_log_probs=behaviour_actions_log_probs,
          target_actions_log_probs=target_actions_log_probs,
          policy_entropy=policy_entropy,
          dones=dones,
          discount=self.gamma,
          rewards=rewards,
          values=values,
          bootstrap_value=bootstrap_value,
          entropy_coeff=entropy_coeff,
          vf_loss_coeff=self.vf_loss_coeff,
          clip_rho_threshold=self.clip_rho_threshold,
          clip_pg_rho_threshold=self.clip_pg_rho_threshold,
      )
  
      self.optimizer.set_lr(learning_rate)
      vtrace_loss.total_loss.backward()
      self.optimizer.step()
      self.optimizer.clear_grad()
  
      return vtrace_loss, kl
  
file_path: PARL/parl/algorithms/paddle/impala/impala.py
imports: import parl
  import paddle
  import paddle.nn.functional
  from paddle.distribution import Categorical
  from parl.algorithms.paddle.impala import vtrace
label_desc: Use image node_id 3 for calling the IMPALA obj's underlying member method code with example usage: obj.learn(obs, actions, behaviour_logits, rewards, dones, learning_rate, entropy_coeff) and returns: vtrace_loss,  kl, paddle