code: def train_function(n_epochs: int, distributed_env_cls) -> A2C:
      """Trains the cart pole environment with the multi-step a2c implementation."""
  
      # initialize distributed env
      envs = distributed_env_cls(
          [lambda: GymMazeEnv(env="CartPole-v0") for _ in range(2)]
      )
  
      # initialize the env and enable statistics collection
      eval_env = distributed_env_cls(
          [lambda: GymMazeEnv(env="CartPole-v0") for _ in range(2)],
          logging_prefix="eval",
      )
  
      # init distribution mapper
      env = GymMazeEnv(env="CartPole-v0")
      distribution_mapper = DistributionMapper(
          action_space=env.action_space, distribution_mapper_config={}
      )
  
      # initialize policies
      policies = {
          0: FlattenConcatPolicyNet(
              {"observation": (4,)},
              {"action": (2,)},
              hidden_units=[16],
              non_lin=nn.Tanh,
          )
      }
  
      # initialize critic
      critics = {
          0: FlattenConcatStateValueNet(
              {"observation": (4,)}, hidden_units=[16], non_lin=nn.Tanh
          )
      }
  
      # algorithm configuration
      algorithm_config = A2CAlgorithmConfig(
          n_epochs=n_epochs,
          epoch_length=2,
          patience=10,
          critic_burn_in_epochs=0,
          n_rollout_steps=20,
          lr=0.0005,
          gamma=0.98,
          gae_lambda=1.0,
          policy_loss_coef=1.0,
          value_loss_coef=0.5,
          entropy_coef=0.0,
          max_grad_norm=0.0,
          device="cpu",
          rollout_evaluator=RolloutEvaluator(
              eval_env=eval_env,
              n_episodes=1,
              model_selection=None,
              deterministic=True,
          ),
      )
  
      # initialize actor critic model
      model = TorchActorCritic(
          policy=TorchPolicy(
              networks=policies,
              distribution_mapper=distribution_mapper,
              device=algorithm_config.device,
          ),
          critic=TorchSharedStateCritic(
              networks=critics,
              obs_spaces_dict=env.observation_spaces_dict,
              device=algorithm_config.device,
              stack_observations=False,
          ),
          device=algorithm_config.device,
      )
  
      a2c = A2C(
          rollout_generator=RolloutGenerator(envs),
          algorithm_config=algorithm_config,
          evaluator=algorithm_config.rollout_evaluator,
          model=model,
          model_selection=None,
      )
  
      # train agent
      a2c.train()
  
      return a2c
  
file_path: maze/maze/test/train/a2c/test_a2c_multi_step.py
imports: import torch.nn
  from maze.core.agent.torch_actor_critic import TorchActorCritic
  from maze.core.agent.torch_policy import TorchPolicy
  from maze.core.agent.torch_state_critic import TorchSharedStateCritic
  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.models.built_in.flatten_concat import FlattenConcatPolicyNet, FlattenConcatStateValueNet
  from maze.train.parallelization.vector_env.sequential_vector_env import SequentialVectorEnv
  from maze.train.parallelization.vector_env.subproc_vector_env import SubprocVectorEnv
  from maze.train.trainers.a2c.a2c_algorithm_config import A2CAlgorithmConfig
  from maze.train.trainers.a2c.a2c_trainer import A2C
  from maze.train.trainers.common.evaluators.rollout_evaluator import RolloutEvaluator
label_desc: Use image node_id 1 for calling a global function with example usage: train_function(n_epochs, distributed_env_cls) and returns: a2c