docstr: input:
      batch: dict(o, s, u, r, u_onehot, avail_u, padded, isover, actor_inputs, critic_inputs)
      q_targets: Q value of target critic network, shape (ep_num, tr_num, n_agents)
  output:
      lambda_return: TD lambda return, shape (ep_num, tr_num, n_agents)
docstr_tok: ["input",":","batch",":","dict","(","o",",","s",",","u",",","r",",","u_onehot",",","avail_u",",","padded",",","isover",",","actor_inputs",",","critic_inputs",")","q_targets",":","Q","value","of","target","critic","network",",","shape","(","ep_num",",","tr_num",",","n_agents",")","output",":","lambda_return",":","TD","lambda","return",",","shape","(","ep_num",",","tr_num",",","n_agents",")"]
code: def _cal_td_target(self, batch, q_targets):  # compute TD(lambda)
      """input:
          batch: dict(o, s, u, r, u_onehot, avail_u, padded, isover, actor_inputs, critic_inputs)
          q_targets: Q value of target critic network, shape (ep_num, tr_num, n_agents)
      output:
          lambda_return: TD lambda return, shape (ep_num, tr_num, n_agents)
      """
      ep_num = batch["r"].shape[0]
      tr_num = batch["r"].shape[1]
      mask = (
          (1 - batch["padded"].float())
          .repeat(1, 1, self.n_agents)
          .to(self.device)
      )
      isover = (
          (1 - batch["isover"].float())
          .repeat(1, 1, self.n_agents)
          .to(self.device)
      )  # used for setting last transition's q_target to 0
      # reshape reward: from (ep_num, tr_num, 1) to (ep_num, tr_num, n_agents)
      r = batch["r"].repeat((1, 1, self.n_agents)).to(self.device)
      # compute n_step_return
      n_step_return = torch.zeros(
          (ep_num, tr_num, self.n_agents, tr_num)
      ).to(self.device)
      for tr_id in range(tr_num - 1, -1, -1):
          n_step_return[:, tr_id, :, 0] = (
              r[:, tr_id]
              + self.gamma * q_targets[:, tr_id] * isover[:, tr_id]
          ) * mask[:, tr_id]
          for n in range(1, tr_num - tr_id):
              n_step_return[:, tr_id, :, n] = (
                  r[:, tr_id]
                  + self.gamma * n_step_return[:, tr_id + 1, :, n - 1]
              ) * mask[:, tr_id]
  
      lambda_return = torch.zeros((ep_num, tr_num, self.n_agents)).to(
          self.device
      )
      for tr_id in range(tr_num):
          returns = torch.zeros((ep_num, self.n_agents)).to(self.device)
          for n in range(1, tr_num - tr_id):
              returns += (
                  pow(self.td_lambda, n - 1)
                  * n_step_return[:, tr_id, :, n - 1]
              )
          lambda_return[:, tr_id] = (
              1 - self.td_lambda
          ) * returns + pow(
              self.td_lambda, tr_num - tr_id - 1
          ) * n_step_return[
              :, tr_id, :, tr_num - tr_id - 1
          ]
      return lambda_return
  
file_path: PARL/parl/algorithms/torch/coma.py
imports: import torch
  import os
  from copy import deepcopy
  import parl
  from parl.utils.utils import check_model_method
  import numpy
label_desc: Use image node_id 6 for calling the COMA obj's underlying member method code with example usage: obj._cal_td_target(batch, q_targets) and returns: lambda_return