class_docstr: Simple feed forward policy network.
      
class_docstr_tok: ["Simple","feed","forward","policy","network","."]
code: def __init__(
      self,
      obs_shapes: Dict[str, Sequence[int]],
      action_logits_shapes: Dict[str, Sequence[int]],
      non_lin: Union[str, type(nn.Module)],
  ):
      super().__init__()
      self.obs_shapes = obs_shapes
      action_key = list(action_logits_shapes.keys())[0]
      # build perception part
      self.perception_dict = OrderedDict()
      self.perception_dict["embedding"] = DenseBlock(
          in_keys="observation",
          out_keys="embedding",
          in_shapes=obs_shapes["observation"],
          hidden_units=[256, 256],
          non_lin=non_lin,
      )
  
      # build action head
      self.perception_dict[action_key] = LinearOutputBlock(
          in_keys="embedding",
          out_keys=action_key,
          in_shapes=self.perception_dict["embedding"].out_shapes(),
          output_units=action_logits_shapes[action_key][0],
      )
  
      self.perception_net = InferenceBlock(
          in_keys="observation",
          out_keys=action_key,
          in_shapes=[self.obs_shapes["observation"]],
          perception_blocks=self.perception_dict,
      )
  
      # initialize model weights
      self.perception_net.apply(make_module_init_normc(1.0))
      self.perception_dict[action_key].apply(
          make_module_init_normc(0.01)
      )
  
imports: from collections import OrderedDict
  from typing import Dict, Union, Sequence
  import torch.nn
  from gym import spaces
  from maze.core.agent.torch_actor_critic import TorchActorCritic
  from maze.core.agent.torch_policy import TorchPolicy
  from maze.core.agent.torch_state_action_critic import TorchStepStateActionCritic
  from maze.core.wrappers.maze_gym_env_wrapper import GymMazeEnv
  from maze.distributions.distribution_mapper import DistributionMapper
  from maze.perception.blocks.feed_forward.dense import DenseBlock
  from maze.perception.blocks.general.concat import ConcatenationBlock
  from maze.perception.blocks.inference import InferenceBlock
  from maze.perception.blocks.output.linear import LinearOutputBlock
  from maze.perception.perception_utils import observation_spaces_to_in_shapes
  from maze.perception.weight_init import make_module_init_normc
  from maze.train.parallelization.vector_env.sequential_vector_env import SequentialVectorEnv
  from maze.train.trainers.common.evaluators.rollout_evaluator import RolloutEvaluator
  from maze.train.parallelization.distributed_actors.dummy_distributed_workers_with_buffer import DummyDistributedWorkersWithBuffer
  from maze.train.trainers.common.replay_buffer.uniform_replay_buffer import UniformReplayBuffer
  from maze.train.trainers.sac.sac_algorithm_config import SACAlgorithmConfig
  from maze.train.trainers.sac.sac_runners import SACRunner
  from maze.train.trainers.sac.sac_trainer import SAC
label_desc: Use image node_id 1 to create a new PolicyNet object from inherited base classes: nn with example: obj = PolicyNet(obs_shapes, action_logits_shapes, non_lin)