class_docstr: Loss function for behavioral cloning.
docstr: Iterate over the action space of a given policy and calculate the loss based on the types of the
  subspaces.
  
  :param actor_id: The actor id to compute the loss for the given substep.
  :param logits_dict: Logits dict output by the policy for a given substep.
  :param target_dict: The policy target for the given substep.
  :param target_actions_dict: The policy target actions for the given substep.
  :param action_spaces_dict: Dict action space for the given substep
  :return: Total loss for this substep
docstr_tok: ["Iterate","over","the","action","space","of","a","given","policy","and","calculate","the","loss","based","on","the","types","of","the","subspaces",".",":","param","actor_id",":","The","actor","id","to","compute","the","loss","for","the","given","substep",".",":","param","logits_dict",":","Logits","dict","output","by","the","policy","for","a","given","substep",".",":","param","target_dict",":","The","policy","target","for","the","given","substep",".",":","param","target_actions_dict",":","The","policy","target","actions","for","the","given","substep",".",":","param","action_spaces_dict",":","Dict","action","space","for","the","given","substep",":","return",":","Total","loss","for","this","substep"]
code: def _get_substep_loss(
      self,
      actor_id: ActorID,
      logits_dict: Dict[str, torch.Tensor],
      target_dict: Dict[str, torch.Tensor],
      target_actions_dict: Dict[str, torch.Tensor],
      action_spaces_dict: gym.spaces.Dict,
      events: ImitationEvents,
  ) -> torch.Tensor:
      """Iterate over the action space of a given policy and calculate the loss based on the types of the
      subspaces.
  
      :param actor_id: The actor id to compute the loss for the given substep.
      :param logits_dict: Logits dict output by the policy for a given substep.
      :param target_dict: The policy target for the given substep.
      :param target_actions_dict: The policy target actions for the given substep.
      :param action_spaces_dict: Dict action space for the given substep
      :return: Total loss for this substep
      """
      losses = []
      for subspace_id, logits in logits_dict.items():
          target = target_dict[subspace_id]
          target_actions = target_actions_dict[subspace_id]
          action_space = action_spaces_dict[subspace_id]
  
          # Categorical (discrete spaces)
          if isinstance(action_space, gym.spaces.Discrete):
              if (
                  logits.dim() == target.dim()
                  and target is target_actions
              ):
                  logits = logits.unsqueeze(0)
              losses.append(self.loss_discrete(logits, target))
  
              events.discrete_accuracy(
                  step_id=actor_id.step_key,
                  agent_id=actor_id.agent_id,
                  subspace_name=subspace_id,
                  value=torch.eq(logits.argmax(dim=-1), target_actions)
                  .float()
                  .mean()
                  .item(),
              )
  
              if logits.shape[-1] > 10:
                  events.discrete_top_5_accuracy(
                      step_id=actor_id.step_key,
                      agent_id=actor_id.agent_id,
                      subspace_name=subspace_id,
                      value=sum(
                          [
                              torch.eq(
                                  logits.argsort(-1)[:, -i],
                                  target_actions,
                              ).float()
                              for i in range(
                                  1, min(5 + 1, logits.shape[-1])
                              )
                          ]
                      )
                      .float()
                      .mean()
                      .item(),
                  )
  
              if logits.shape[-1] > 20:
                  events.discrete_top_10_accuracy(
                      step_id=actor_id.step_key,
                      agent_id=actor_id.agent_id,
                      subspace_name=subspace_id,
                      value=sum(
                          [
                              torch.eq(
                                  logits.argsort(-1)[:, -i],
                                  target_actions,
                              ).float()
                              for i in range(
                                  1, min(10 + 1, logits.shape[-1])
                              )
                          ]
                      )
                      .float()
                      .mean()
                      .item(),
                  )
  
              batch_ranks = torch.where(
                  torch.argsort(logits, dim=-1, descending=True)
                  == target_actions.unsqueeze(-1)
              )[1]
              events.discrete_action_rank(
                  step_id=actor_id.step_key,
                  agent_id=actor_id.agent_id,
                  subspace_name=subspace_id,
                  value=batch_ranks.float().mean().item(),
              )
  
          # Multi-binary (multi-binary spaces)
          elif isinstance(action_space, gym.spaces.MultiBinary):
              losses.append(self.loss_multi_binary(logits, target))
              events.multi_binary_accuracy(
                  step_id=actor_id.step_key,
                  agent_id=actor_id.agent_id,
                  subspace_name=subspace_id,
                  value=torch.eq(logits > 0, target_actions)
                  .float()
                  .mean()
                  .item(),
              )
  
          # Continuous (box spaces)
          elif isinstance(action_space, gym.spaces.Box):
              pred, _ = torch.chunk(logits, chunks=2, dim=-1)
              losses.append(self.loss_box(pred.float(), target.float()))
              events.box_mean_abs_deviation(
                  step_id=actor_id.step_key,
                  agent_id=actor_id.agent_id,
                  subspace_name=subspace_id,
                  value=target.sub(pred).abs().mean().item(),
              )
  
          else:
              raise NotImplementedError(
                  "Only Discrete, Box, and MultiBinary action spaces are supported."
              )
  
      loss = sum(losses)
      events.policy_loss(
          step_id=actor_id.step_key,
          agent_id=actor_id.agent_id,
          value=loss.item(),
      )
      return loss
  
file_path: maze/maze/train/trainers/imitation/bc_loss.py
imports: from dataclasses import dataclass
  from typing import Dict, Union, List, Optional
  import gym
  import torch
  import torch.nn
  import torch.nn.functional
  from maze.core.agent.torch_policy import TorchPolicy
  from maze.core.env.action_conversion import TorchActionType
  from maze.core.env.observation_conversion import ObservationType
  from maze.core.env.structured_env import ActorID
  from maze.train.trainers.imitation.imitation_events import ImitationEvents
label_desc: Use image node_id 2 for calling the BCLoss obj's underlying member method code with example usage: obj._get_substep_loss(actor_id, logits_dict, target_dict, target_actions_dict, action_spaces_dict, events) and returns: loss