code: def build_structured_with_critic_type(
      env,
      critics_composer_type: type(BaseStateCriticComposer),
      critics_type: type(TorchStateCritic),
      shared_embedding_keys: Optional[
          Union[List[str], Dict[StepKeyType, List[str]]]
      ],
  ):
      """helper function"""
  
      # map observations to a modality
      obs_modalities = {
          "observation_0": "image",
          "observation_1": "feature",
          DeltaStateCriticComposer.prev_value_key: "feature",
      }
  
      # define how to process a modality
      modality_config = dict()
      modality_config["feature"] = {
          "block_type": "maze.perception.blocks.DenseBlock",
          "block_params": {
              "hidden_units": [32, 32],
              "non_lin": "torch.nn.ReLU",
          },
      }
      modality_config["image"] = {
          "block_type": "maze.perception.blocks.StridedConvolutionDenseBlock",
          "block_params": {
              "hidden_channels": [8, 16, 32],
              "hidden_kernels": [8, 4, 4],
              "convolution_dimension": 2,
              "hidden_strides": [4, 2, 2],
              "hidden_dilations": None,
              "hidden_padding": [1, 1, 1],
              "padding_mode": None,
              "hidden_units": [],
              "non_lin": "torch.nn.SELU",
          },
      }
  
      modality_config["hidden"] = {
          "block_type": "maze.perception.blocks.DenseBlock",
          "block_params": {
              "hidden_units": [64],
              "non_lin": "torch.nn.ReLU",
          },
      }
      modality_config["recurrence"] = {}
  
      model_builder = {
          "_target_": "maze.perception.builders.concat.ConcatModelBuilder",
          "modality_config": modality_config,
          "observation_modality_mapping": obs_modalities,
          "shared_embedding_keys": shared_embedding_keys,
      }
  
      # initialize default model builder
      default_builder = TemplateModelComposer(
          action_spaces_dict=env.action_spaces_dict,
          observation_spaces_dict=env.observation_spaces_dict,
          agent_counts_dict=env.agent_counts_dict,
          distribution_mapper_config={},
          model_builder=model_builder,
          policy={"_target_": policy_composer_type},
          critic={"_target_": critics_composer_type},
      )
  
      # create model pdf
      default_builder.save_models()
  
      assert isinstance(
          default_builder.distribution_mapper, DistributionMapper
      )
      for pp in default_builder.policy.networks.values():
          assert isinstance(pp, nn.Module)
      for cc in default_builder.critic.networks.values():
          assert isinstance(cc, nn.Module)
  
      assert isinstance(default_builder.critic, critics_type)
  
      rollout_generator = RolloutGenerator(
          env=env, record_next_observations=False
      )
      policy = RandomPolicy(env.action_spaces_dict)
      trajectory = (
          rollout_generator.rollout(policy, n_steps=10)
          .stack()
          .to_torch(device="cpu")
      )
  
      policy_output = default_builder.policy.compute_policy_output(
          trajectory
      )
      critic_input = StateCriticInput.build(policy_output, trajectory)
      _ = default_builder.critic.predict_values(critic_input)
  
file_path: maze/maze/test/perception/models/test_template_model_composer.py
imports: from typing import Dict, List, Union, Optional
  import pytest
  from torch import nn
  from maze.core.agent.random_policy import RandomPolicy
  from maze.core.agent.state_critic_input_output import StateCriticInput
  from maze.core.agent.torch_state_critic import TorchSharedStateCritic, TorchDeltaStateCritic, TorchStepStateCritic, TorchStateCritic
  from maze.core.env.structured_env import StepKeyType
  from maze.core.rollout.rollout_generator import RolloutGenerator
  from maze.core.wrappers.maze_gym_env_wrapper import GymMazeEnv
  from maze.distributions.distribution_mapper import DistributionMapper
  from maze.perception.blocks.inference import InferenceBlock
  from maze.perception.models.critics import SharedStateCriticComposer, StepStateCriticComposer, BaseStateCriticComposer, DeltaStateCriticComposer
  from maze.perception.models.template_model_composer import TemplateModelComposer
  from maze.test.shared_test_utils.helper_functions import build_dummy_structured_env, build_dummy_maze_env_with_structured_core_env
label_desc: Use image node_id 3 for calling a global function with example usage: build_structured_with_critic_type(env, critics_composer_type, critics_type, shared_embedding_keys) without return types