class_docstr: The Policy net (actor) computing the action probabilities from the observations.
  
  :param obs_shapes: The shapes of all observations as a dict.
  :param action_logits_shapes: The shapes of all actions as a dict structure.
  :param non_lin: The nonlinear activation to be used.
  :param with_mask: Weather to use action masking or not.
class_docstr_tok: ["The","Policy","net","(","actor",")","computing","the","action","probabilities","from","the","observations",".",":","param","obs_shapes",":","The","shapes","of","all","observations","as","a","dict",".",":","param","action_logits_shapes",":","The","shapes","of","all","actions","as","a","dict","structure",".",":","param","non_lin",":","The","nonlinear","activation","to","be","used",".",":","param","with_mask",":","Weather","to","use","action","masking","or","not","."]
code: def __init__(
      self,
      obs_shapes: Dict[str, Sequence[int]],
      action_logits_shapes: Dict[str, Sequence[int]],
      non_lin: Union[str, type(nn.Module)],
      with_mask: bool,
  ):
      nn.Module.__init__(self)
      self.obs_shapes = obs_shapes
  
      hidden_units = 32
  
      self.perception_dict = OrderedDict()
  
      self.perception_dict["selected_feat"] = DenseBlock(
          in_keys="selected_piece",
          out_keys="selected_feat",
          in_shapes=self.obs_shapes["selected_piece"],
          hidden_units=[hidden_units],
          non_lin=non_lin,
      )
  
      self.perception_dict["order_feat"] = DenseBlock(
          in_keys="ordered_piece",
          out_keys="order_feat",
          in_shapes=self.obs_shapes["ordered_piece"],
          hidden_units=[hidden_units],
          non_lin=non_lin,
      )
  
      self.perception_dict["latent"] = ConcatenationBlock(
          in_keys=["selected_feat", "order_feat"],
          out_keys="latent",
          in_shapes=[[hidden_units], [hidden_units]],
          concat_dim=-1,
      )
  
      rotation_out_key = (
          "cut_rotation_logits" if with_mask else "cut_rotation"
      )
      self.perception_dict[rotation_out_key] = LinearOutputBlock(
          in_keys="latent",
          out_keys=rotation_out_key,
          in_shapes=self.perception_dict["latent"].out_shapes(),
          output_units=action_logits_shapes["cut_rotation"][0],
      )
  
      if with_mask:
          self.perception_dict["cut_rotation"] = ActionMaskingBlock(
              in_keys=["cut_rotation_logits", "cutting_mask"],
              out_keys="cut_rotation",
              in_shapes=self.perception_dict[
                  "cut_rotation_logits"
              ].out_shapes()
              + [self.obs_shapes["cutting_mask"]],
              num_actors=1,
              num_of_actor_actions=None,
          )
  
      self.perception_dict["cut_order"] = LinearOutputBlock(
          in_keys="latent",
          out_keys="cut_order",
          in_shapes=self.perception_dict["latent"].out_shapes(),
          output_units=action_logits_shapes["cut_order"][0],
      )
  
      in_keys = ["selected_piece", "ordered_piece"]
      if with_mask:
          in_keys.append("cutting_mask")
      self.perception_net = InferenceBlock(
          in_keys=in_keys,
          out_keys=["cut_rotation", "cut_order"],
          in_shapes=[self.obs_shapes[key] for key in in_keys],
          perception_blocks=self.perception_dict,
      )
  
      # initialize model weights
      self.perception_net.apply(make_module_init_normc(1.0))
      self.perception_dict[rotation_out_key].apply(
          make_module_init_normc(0.01)
      )
      self.perception_dict["cut_order"].apply(
          make_module_init_normc(0.01)
      )
  
imports: from collections import OrderedDict
  from typing import Dict, Union, Sequence
  import torch
  from maze.perception.blocks.feed_forward.dense import DenseBlock
  from maze.perception.blocks.general.action_masking import ActionMaskingBlock
  from maze.perception.blocks.general.concat import ConcatenationBlock
  from maze.perception.blocks.general.correlation import CorrelationBlock
  from maze.perception.blocks.general.functional import FunctionalBlock
  from maze.perception.blocks.inference import InferenceBlock
  from maze.perception.blocks.output.linear import LinearOutputBlock
  from maze.perception.weight_init import make_module_init_normc
  from torch import nn
label_desc: Use image node_id 1 to create a new CuttingPolicyNet object from inherited base classes: nn with example: obj = CuttingPolicyNet(obs_shapes, action_logits_shapes, non_lin, with_mask)