class_docstr: Encapsulates multiple torch state critics for training in structured environments.
  
  :param networks: Mapping of value functions (critic) to encapsulate.
  :param obs_spaces_dict: The observation spaces dict of the environment.
  :param device: Device the policy should be located on (cpu or cuda).
class_docstr_tok: ["Encapsulates","multiple","torch","state","critics","for","training","in","structured","environments",".",":","param","networks",":","Mapping","of","value","functions","(","critic",")","to","encapsulate",".",":","param","obs_spaces_dict",":","The","observation","spaces","dict","of","the","environment",".",":","param","device",":","Device","the","policy","should","be","located","on","(","cpu","or","cuda",")","."]
docstr: Compute bootstrapped return from rewards and estimated values.
  
  :param gamma: Discounting factor
  :param gae_lambda: Bias vs variance trade of factor for Generalized Advantage Estimator (GAE)
  :param rewards: Step rewards with shape (n_steps, n_workers)
  :param values: Predicted values with shape (n_steps, n_workers)
  :param dones: Step dones with shape (n_steps, n_workers)
  :param deltas: Predicted value deltas to previous sub-step with shape (n_steps, n_workers)
  :return: Per time step returns.
docstr_tok: ["Compute","bootstrapped","return","from","rewards","and","estimated","values",".",":","param","gamma",":","Discounting","factor",":","param","gae_lambda",":","Bias","vs","variance","trade","of","factor","for","Generalized","Advantage","Estimator","(","GAE",")",":","param","rewards",":","Step","rewards","with","shape","(","n_steps",",","n_workers",")",":","param","values",":","Predicted","values","with","shape","(","n_steps",",","n_workers",")",":","param","dones",":","Step","dones","with","shape","(","n_steps",",","n_workers",")",":","param","deltas",":","Predicted","value","deltas","to","previous","sub-step","with","shape","(","n_steps",",","n_workers",")",":","return",":","Per","time","step","returns","."]
code: def compute_return(
      self,
      gamma: float,
      gae_lambda: float,
      rewards: torch.Tensor,
      values: torch.Tensor,
      dones: torch.Tensor,
      deltas: torch.Tensor = None,
  ) -> torch.Tensor:
      """Compute bootstrapped return from rewards and estimated values.
  
      :param gamma: Discounting factor
      :param gae_lambda: Bias vs variance trade of factor for Generalized Advantage Estimator (GAE)
      :param rewards: Step rewards with shape (n_steps, n_workers)
      :param values: Predicted values with shape (n_steps, n_workers)
      :param dones: Step dones with shape (n_steps, n_workers)
      :param deltas: Predicted value deltas to previous sub-step with shape (n_steps, n_workers)
      :return: Per time step returns.
      """
      assert (
          rewards.shape == values.shape
      ), f"{rewards.shape} vs {values.shape}"
      assert (
          rewards.shape == dones.shape
      ), f"{rewards.shape} vs {dones.shape}"
  
      # initialize returns
      returns = torch.zeros(
          (rewards.shape[0], rewards.shape[1]),
          dtype=torch.float32,
          device=self.device,
      )
  
      # prepare end-of-episode mask
      mask = (~dones).float()
  
      # traverse time steps in reverse order
      gae = torch.zeros(
          rewards.shape[1], dtype=torch.float32, device=self.device
      )
      for t in reversed(range(0, len(rewards))):
          # bootstrap value function for last entry
          if t == len(rewards) - 1:
              returns[t] = values[t]
              if deltas is not None:
                  returns[t] = values[t] + deltas[t]
  
          # compute discounted return
          else:
              if gae_lambda != 1.0:
                  delta = (
                      rewards[t]
                      + gamma * values[t + 1] * mask[t]
                      - values[t]
                  )
                  gae = delta + gamma * gae_lambda * gae
                  returns[t] = gae + values[t]
              else:
                  returns[t] = (
                      rewards[t] + gamma * returns[t + 1] * mask[t]
                  )
  
      return returns
  
imports: from abc import abstractmethod
  from typing import Mapping, Union, List, Dict
  import torch
  from gym import spaces
  from torch import nn
  from maze.core.agent.state_critic import StateCritic
  from maze.core.agent.state_critic_input_output import StateCriticStepOutput, StateCriticOutput, StateCriticInput
  from maze.core.agent.torch_model import TorchModel
  from maze.core.annotations import override
  from maze.core.env.observation_conversion import ObservationType
  from maze.core.env.structured_env import StepKeyType
  from maze.perception.perception_utils import convert_to_torch, flatten_spaces, stack_and_flatten_spaces
label_desc: Use image node_id 12 for calling the TorchStateCritic obj's underlying member method code with example usage: obj.compute_return(gamma, gae_lambda, rewards, values, dones, deltas) and returns: returns