code: def test_cartpole_policy_model():
      flat_action_space = flat_structured_space(
          _mock_action_spaces_dict()
      )
      distribution_mapper = DistributionMapper(
          action_space=flat_action_space, distribution_mapper_config={}
      )
  
      action_logits_shapes = {
          step_key: {
              action_head: distribution_mapper.required_logits_shape(
                  action_head
              )
              for action_head in _mock_action_spaces_dict()[
                  step_key
              ].spaces.keys()
          }
          for step_key in _mock_action_spaces_dict().keys()
      }
  
      obs_shapes = observation_spaces_to_in_shapes(
          _mock_observation_spaces_dict()
      )
  
      policy = CustomComplexPolicyNet(
          obs_shapes[0],
          action_logits_shapes[0],
          non_lin="torch.nn.ReLU",
          hidden_units=[128],
      )
  
      critic = CustomComplexCriticNet(
          obs_shapes[0], non_lin="torch.nn.ReLU", hidden_units=[128]
      )
  
      obs_np = _mock_observation_spaces_dict()[0].sample()
      obs = {k: torch.from_numpy(v) for k, v in obs_np.items()}
  
      actions = policy(obs)
      values = critic(obs)
  
      assert "action_move" in actions
      assert "action_use" in actions
      assert "value" in values
  
file_path: maze/docs/test/policy_and_value_networks/test_complex_models.py
imports: import torch
  import yaml
  from gym import spaces
  import docs.source.policy_and_value_networks.code_snippets
  from docs.source.policy_and_value_networks.code_snippets.custom_complex_critic_net import CustomComplexCriticNet
  from docs.source.policy_and_value_networks.code_snippets.custom_complex_policy_net import CustomComplexPolicyNet
  from maze.core.utils.factory import Factory
  from maze.core.utils.structured_env_utils import flat_structured_space
  from maze.distributions.distribution_mapper import DistributionMapper
  from maze.perception.models.model_composer import BaseModelComposer
  from maze.perception.perception_utils import observation_spaces_to_in_shapes
label_desc: Use image node_id 4 for calling a global function with example usage: test_cartpole_policy_model() without return types