fixing errors

ml-agents-envs/mlagents_envs/base_env.py

     HYBRID = 2
 
 
-class BehaviorSpec(NamedTuple):
+class HybridBehaviorSpec(NamedTuple):
     observation_shapes: List[Tuple]
     continuous_action_shape: int
     discrete_action_shape: Tuple[int]

ml-agents/mlagents/trainers/policy/policy.py

         self.network_settings: NetworkSettings = trainer_settings.network_settings
         self.seed = seed
         # For hybrid
-        self.continuous_act_size = behavior_spec.continuous_action_size()
-        self.discrete_act_size = behavior_spec.continuous_action_size()
-        self.act_size = (
-            list(behavior_spec.discrete_action_branches)
-            if behavior_spec.is_action_discrete()
-            else [behavior_spec.action_size]
-        )
+        self.continuous_act_size = behavior_spec.continuous_action_size
+        self.discrete_act_size = behavior_spec.discrete_action_branches
+        #self.act_size = (
+        #    list(behavior_spec.discrete_action_branches)
+        #    if behavior_spec.is_action_discrete()
+        #    else [behavior_spec.action_size]
+        #)
         self.vec_obs_size = sum(
             shape[0] for shape in behavior_spec.observation_shapes if len(shape) == 1
         )
-        self.use_continuous_act = behavior_spec.is_action_continuous()
+        #self.use_continuous_act = behavior_spec.is_action_continuous()
         self.num_branches = self.behavior_spec.action_size
         self.previous_action_dict: Dict[str, np.array] = {}
         self.memory_dict: Dict[str, np.ndarray] = {}

ml-agents/mlagents/trainers/policy/torch_policy.py

 
 from mlagents.trainers.settings import TrainerSettings
 from mlagents.trainers.trajectory import SplitObservations
+from mlagents.trainers.torch.distributions import DistInstance
 from mlagents.trainers.torch.networks import (
     SharedActorCritic,
     SeparateActorCritic,
         self.actor_critic = ac_class(
             observation_shapes=self.behavior_spec.observation_shapes,
             network_settings=trainer_settings.network_settings,
-            act_type=behavior_spec.action_type,
             continuous_act_size=self.continuous_act_size,
             discrete_act_size=self.discrete_act_size,
             stream_names=reward_signal_names,
         if self.use_vec_obs and self.normalize:
             self.actor_critic.update_normalization(vector_obs)
 
-    def get_actions_and_stats(dists : List[DistInstance]):
-    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    def get_actions_and_stats(dists : List[DistInstance]) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    
         action_list = self.actor_critic.sample_action(dists)
         log_probs, entropies, all_logs = ModelUtils.get_probs_and_entropy(
             action_list, dists

ml-agents/mlagents/trainers/tests/simple_test_envs.py

 from mlagents_envs.base_env import (
     BaseEnv,
     BehaviorSpec,
-    DecisionSteps,
+    DecisionSteps,
     TerminalSteps,
     ActionType,
     BehaviorMapping,
             vis_obs_size,
             vec_obs_size,
             action_size,
+        )
+        super().__init__(
+            brain_names,
+            False,
+            step_size=step_size,
+            num_visual=num_visual,
+            num_vector=num_vector,
         )
         # Number of steps to reveal the goal for. Lower is harder. Should be
         # less than 1/step_size to force agent to use memory

ml-agents/mlagents/trainers/tests/torch/test_simple_rl.py

 
 from mlagents.trainers.tests.simple_test_envs import (
     SimpleEnvironment,
+    HybridEnvironment,
     MemoryEnvironment,
     RecordEnvironment,
 )
             assert all(reward > success_threshold for reward in processed_rewards)
 
 
-@pytest.mark.parametrize("use_discrete", [True, False])
-def test_simple_ppo(use_discrete):
-    env = SimpleEnvironment([BRAIN_NAME], use_discrete=use_discrete)
+#@pytest.mark.parametrize("use_discrete", [True, False])
+#def test_simple_ppo(use_discrete):
+#    env = SimpleEnvironment([BRAIN_NAME], use_discrete=use_discrete)
+#    config = attr.evolve(PPO_CONFIG)
+#    _check_environment_trains(env, {BRAIN_NAME: config})
+#
+def test_hybrid_ppo():
+    env = HybridEnvironment([BRAIN_NAME])
-@pytest.mark.parametrize("use_discrete", [True, False])
-def test_2d_ppo(use_discrete):
-    env = SimpleEnvironment(
-        [BRAIN_NAME], use_discrete=use_discrete, action_size=2, step_size=0.8
-    )
-    new_hyperparams = attr.evolve(
-        PPO_CONFIG.hyperparameters, batch_size=64, buffer_size=640
-    )
-    config = attr.evolve(PPO_CONFIG, hyperparameters=new_hyperparams, max_steps=10000)
-    _check_environment_trains(env, {BRAIN_NAME: config})
-
-
-@pytest.mark.parametrize("use_discrete", [True, False])
-@pytest.mark.parametrize("num_visual", [1, 2])
-def test_visual_ppo(num_visual, use_discrete):
-    env = SimpleEnvironment(
-        [BRAIN_NAME],
-        use_discrete=use_discrete,
-        num_visual=num_visual,
-        num_vector=0,
-        step_size=0.2,
-    )
-    new_hyperparams = attr.evolve(PPO_CONFIG.hyperparameters, learning_rate=3.0e-4)
-    config = attr.evolve(PPO_CONFIG, hyperparameters=new_hyperparams)
-    _check_environment_trains(env, {BRAIN_NAME: config})
-
-
-@pytest.mark.parametrize("num_visual", [1, 2])
-@pytest.mark.parametrize("vis_encode_type", ["resnet", "nature_cnn", "match3"])
-def test_visual_advanced_ppo(vis_encode_type, num_visual):
-    env = SimpleEnvironment(
-        [BRAIN_NAME],
-        use_discrete=True,
-        num_visual=num_visual,
-        num_vector=0,
-        step_size=0.5,
-        vis_obs_size=(5, 5, 5) if vis_encode_type == "match3" else (36, 36, 3),
-    )
-    new_networksettings = attr.evolve(
-        SAC_CONFIG.network_settings, vis_encode_type=EncoderType(vis_encode_type)
-    )
-    new_hyperparams = attr.evolve(PPO_CONFIG.hyperparameters, learning_rate=3.0e-4)
-    config = attr.evolve(
-        PPO_CONFIG,
-        hyperparameters=new_hyperparams,
-        network_settings=new_networksettings,
-        max_steps=700,
-        summary_freq=100,
-    )
-    # The number of steps is pretty small for these encoders
-    _check_environment_trains(env, {BRAIN_NAME: config}, success_threshold=0.5)
-
-
-@pytest.mark.parametrize("use_discrete", [True, False])
-def test_recurrent_ppo(use_discrete):
-    env = MemoryEnvironment([BRAIN_NAME], use_discrete=use_discrete)
-    new_network_settings = attr.evolve(
-        PPO_CONFIG.network_settings,
-        memory=NetworkSettings.MemorySettings(memory_size=16),
-    )
-    new_hyperparams = attr.evolve(
-        PPO_CONFIG.hyperparameters, learning_rate=1.0e-3, batch_size=64, buffer_size=128
-    )
-    config = attr.evolve(
-        PPO_CONFIG,
-        hyperparameters=new_hyperparams,
-        network_settings=new_network_settings,
-        max_steps=5000,
-    )
-    _check_environment_trains(env, {BRAIN_NAME: config}, success_threshold=0.9)
-
-
-@pytest.mark.parametrize("use_discrete", [True, False])
-def test_simple_sac(use_discrete):
-    env = SimpleEnvironment([BRAIN_NAME], use_discrete=use_discrete)
-    config = attr.evolve(SAC_CONFIG)
-    _check_environment_trains(env, {BRAIN_NAME: config})
-
-
-@pytest.mark.parametrize("use_discrete", [True, False])
-def test_2d_sac(use_discrete):
-    env = SimpleEnvironment(
-        [BRAIN_NAME], use_discrete=use_discrete, action_size=2, step_size=0.8
-    )
-    new_hyperparams = attr.evolve(SAC_CONFIG.hyperparameters, buffer_init_steps=2000)
-    config = attr.evolve(SAC_CONFIG, hyperparameters=new_hyperparams, max_steps=10000)
-    _check_environment_trains(env, {BRAIN_NAME: config}, success_threshold=0.8)
-
-
-@pytest.mark.parametrize("use_discrete", [True, False])
-@pytest.mark.parametrize("num_visual", [1, 2])
-def test_visual_sac(num_visual, use_discrete):
-    env = SimpleEnvironment(
-        [BRAIN_NAME],
-        use_discrete=use_discrete,
-        num_visual=num_visual,
-        num_vector=0,
-        step_size=0.2,
-    )
-    new_hyperparams = attr.evolve(
-        SAC_CONFIG.hyperparameters, batch_size=16, learning_rate=3e-4
-    )
-    config = attr.evolve(SAC_CONFIG, hyperparameters=new_hyperparams)
-    _check_environment_trains(env, {BRAIN_NAME: config})
-
-
-@pytest.mark.parametrize("num_visual", [1, 2])
-@pytest.mark.parametrize("vis_encode_type", ["resnet", "nature_cnn", "match3"])
-def test_visual_advanced_sac(vis_encode_type, num_visual):
-    env = SimpleEnvironment(
-        [BRAIN_NAME],
-        use_discrete=True,
-        num_visual=num_visual,
-        num_vector=0,
-        step_size=0.5,
-        vis_obs_size=(5, 5, 5) if vis_encode_type == "match3" else (36, 36, 3),
-    )
-    new_networksettings = attr.evolve(
-        SAC_CONFIG.network_settings, vis_encode_type=EncoderType(vis_encode_type)
-    )
-    new_hyperparams = attr.evolve(
-        SAC_CONFIG.hyperparameters,
-        batch_size=16,
-        learning_rate=3e-4,
-        buffer_init_steps=0,
-    )
-    config = attr.evolve(
-        SAC_CONFIG,
-        hyperparameters=new_hyperparams,
-        network_settings=new_networksettings,
-        max_steps=100,
-    )
-    # The number of steps is pretty small for these encoders
-    _check_environment_trains(env, {BRAIN_NAME: config}, success_threshold=0.5)
-
-
-@pytest.mark.parametrize("use_discrete", [True, False])
-def test_recurrent_sac(use_discrete):
-    step_size = 0.2 if use_discrete else 0.5
-    env = MemoryEnvironment(
-        [BRAIN_NAME], use_discrete=use_discrete, step_size=step_size
-    )
-    new_networksettings = attr.evolve(
-        SAC_CONFIG.network_settings,
-        memory=NetworkSettings.MemorySettings(memory_size=16, sequence_length=16),
-    )
-    new_hyperparams = attr.evolve(
-        SAC_CONFIG.hyperparameters,
-        batch_size=128,
-        learning_rate=1e-3,
-        buffer_init_steps=1000,
-        steps_per_update=2,
-    )
-    config = attr.evolve(
-        SAC_CONFIG,
-        hyperparameters=new_hyperparams,
-        network_settings=new_networksettings,
-        max_steps=5000,
-    )
-    _check_environment_trains(env, {BRAIN_NAME: config})
-
-
-@pytest.mark.parametrize("use_discrete", [True, False])
-def test_simple_ghost(use_discrete):
-    env = SimpleEnvironment(
-        [BRAIN_NAME + "?team=0", BRAIN_NAME + "?team=1"], use_discrete=use_discrete
-    )
-    self_play_settings = SelfPlaySettings(
-        play_against_latest_model_ratio=1.0, save_steps=2000, swap_steps=2000
-    )
-    config = attr.evolve(PPO_CONFIG, self_play=self_play_settings, max_steps=2500)
-    _check_environment_trains(env, {BRAIN_NAME: config})
-
-
-@pytest.mark.parametrize("use_discrete", [True, False])
-def test_simple_ghost_fails(use_discrete):
-    env = SimpleEnvironment(
-        [BRAIN_NAME + "?team=0", BRAIN_NAME + "?team=1"], use_discrete=use_discrete
-    )
-    # This config should fail because the ghosted policy is never swapped with a competent policy.
-    # Swap occurs after max step is reached.
-    self_play_settings = SelfPlaySettings(
-        play_against_latest_model_ratio=1.0, save_steps=2000, swap_steps=4000
-    )
-    config = attr.evolve(PPO_CONFIG, self_play=self_play_settings, max_steps=2500)
-    _check_environment_trains(env, {BRAIN_NAME: config}, success_threshold=None)
-    processed_rewards = [
-        default_reward_processor(rewards) for rewards in env.final_rewards.values()
-    ]
-    success_threshold = 0.9
-    assert any(reward > success_threshold for reward in processed_rewards) and any(
-        reward < success_threshold for reward in processed_rewards
-    )
-
-
-@pytest.mark.parametrize("use_discrete", [True, False])
-def test_simple_asymm_ghost(use_discrete):
-    # Make opponent for asymmetric case
-    brain_name_opp = BRAIN_NAME + "Opp"
-    env = SimpleEnvironment(
-        [BRAIN_NAME + "?team=0", brain_name_opp + "?team=1"], use_discrete=use_discrete
-    )
-    self_play_settings = SelfPlaySettings(
-        play_against_latest_model_ratio=1.0,
-        save_steps=10000,
-        swap_steps=10000,
-        team_change=400,
-    )
-    config = attr.evolve(PPO_CONFIG, self_play=self_play_settings, max_steps=4000)
-    _check_environment_trains(env, {BRAIN_NAME: config, brain_name_opp: config})
-
-
-@pytest.mark.parametrize("use_discrete", [True, False])
-def test_simple_asymm_ghost_fails(use_discrete):
-    # Make opponent for asymmetric case
-    brain_name_opp = BRAIN_NAME + "Opp"
-    env = SimpleEnvironment(
-        [BRAIN_NAME + "?team=0", brain_name_opp + "?team=1"], use_discrete=use_discrete
-    )
-    # This config should fail because the team that us not learning when both have reached
-    # max step should be executing the initial, untrained poliy.
-    self_play_settings = SelfPlaySettings(
-        play_against_latest_model_ratio=0.0,
-        save_steps=5000,
-        swap_steps=5000,
-        team_change=2000,
-    )
-    config = attr.evolve(PPO_CONFIG, self_play=self_play_settings, max_steps=3000)
-    _check_environment_trains(
-        env, {BRAIN_NAME: config, brain_name_opp: config}, success_threshold=None
-    )
-    processed_rewards = [
-        default_reward_processor(rewards) for rewards in env.final_rewards.values()
-    ]
-    success_threshold = 0.9
-    assert any(reward > success_threshold for reward in processed_rewards) and any(
-        reward < success_threshold for reward in processed_rewards
-    )
-
-
-@pytest.fixture(scope="session")
-def simple_record(tmpdir_factory):
-    def record_demo(use_discrete, num_visual=0, num_vector=1):
-        env = RecordEnvironment(
-            [BRAIN_NAME],
-            use_discrete=use_discrete,
-            num_visual=num_visual,
-            num_vector=num_vector,
-            n_demos=100,
-        )
-        # If we want to use true demos, we can solve the env in the usual way
-        # Otherwise, we can just call solve to execute the optimal policy
-        env.solve()
-        agent_info_protos = env.demonstration_protos[BRAIN_NAME]
-        meta_data_proto = DemonstrationMetaProto()
-        brain_param_proto = BrainParametersProto(
-            vector_action_size=[2] if use_discrete else [1],
-            vector_action_descriptions=[""],
-            vector_action_space_type=discrete if use_discrete else continuous,
-            brain_name=BRAIN_NAME,
-            is_training=True,
-        )
-        action_type = "Discrete" if use_discrete else "Continuous"
-        demo_path_name = "1DTest" + action_type + ".demo"
-        demo_path = str(tmpdir_factory.mktemp("tmp_demo").join(demo_path_name))
-        write_demo(demo_path, meta_data_proto, brain_param_proto, agent_info_protos)
-        return demo_path
-
-    return record_demo
-
-
-@pytest.mark.parametrize("use_discrete", [True, False])
-@pytest.mark.parametrize("trainer_config", [PPO_CONFIG, SAC_CONFIG])
-def test_gail(simple_record, use_discrete, trainer_config):
-    demo_path = simple_record(use_discrete)
-    env = SimpleEnvironment([BRAIN_NAME], use_discrete=use_discrete, step_size=0.2)
-    bc_settings = BehavioralCloningSettings(demo_path=demo_path, steps=1000)
-    reward_signals = {
-        RewardSignalType.GAIL: GAILSettings(encoding_size=32, demo_path=demo_path)
-    }
-    config = attr.evolve(
-        trainer_config,
-        reward_signals=reward_signals,
-        behavioral_cloning=bc_settings,
-        max_steps=500,
-    )
-    _check_environment_trains(env, {BRAIN_NAME: config}, success_threshold=0.9)
-
-
-@pytest.mark.parametrize("use_discrete", [True, False])
-def test_gail_visual_ppo(simple_record, use_discrete):
-    demo_path = simple_record(use_discrete, num_visual=1, num_vector=0)
-    env = SimpleEnvironment(
-        [BRAIN_NAME],
-        num_visual=1,
-        num_vector=0,
-        use_discrete=use_discrete,
-        step_size=0.2,
-    )
-    bc_settings = BehavioralCloningSettings(demo_path=demo_path, steps=1500)
-    reward_signals = {
-        RewardSignalType.GAIL: GAILSettings(encoding_size=32, demo_path=demo_path)
-    }
-    hyperparams = attr.evolve(PPO_CONFIG.hyperparameters, learning_rate=3e-4)
-    config = attr.evolve(
-        PPO_CONFIG,
-        reward_signals=reward_signals,
-        hyperparameters=hyperparams,
-        behavioral_cloning=bc_settings,
-        max_steps=1000,
-    )
-    _check_environment_trains(env, {BRAIN_NAME: config}, success_threshold=0.9)
-
-
-@pytest.mark.parametrize("use_discrete", [True, False])
-def test_gail_visual_sac(simple_record, use_discrete):
-    demo_path = simple_record(use_discrete, num_visual=1, num_vector=0)
-    env = SimpleEnvironment(
-        [BRAIN_NAME],
-        num_visual=1,
-        num_vector=0,
-        use_discrete=use_discrete,
-        step_size=0.2,
-    )
-    bc_settings = BehavioralCloningSettings(demo_path=demo_path, steps=1000)
-    reward_signals = {
-        RewardSignalType.GAIL: GAILSettings(encoding_size=32, demo_path=demo_path)
-    }
-    hyperparams = attr.evolve(
-        SAC_CONFIG.hyperparameters, learning_rate=3e-4, batch_size=16
-    )
-    config = attr.evolve(
-        SAC_CONFIG,
-        reward_signals=reward_signals,
-        hyperparameters=hyperparams,
-        behavioral_cloning=bc_settings,
-        max_steps=500,
-    )
-    _check_environment_trains(env, {BRAIN_NAME: config}, success_threshold=0.9)
+#
+#@pytest.mark.parametrize("use_discrete", [True, False])
+#def test_2d_ppo(use_discrete):
+#    env = SimpleEnvironment(
+#        [BRAIN_NAME], use_discrete=use_discrete, action_size=2, step_size=0.8
+#    )
+#    new_hyperparams = attr.evolve(
+#        PPO_CONFIG.hyperparameters, batch_size=64, buffer_size=640
+#    )
+#    config = attr.evolve(PPO_CONFIG, hyperparameters=new_hyperparams, max_steps=10000)
+#    _check_environment_trains(env, {BRAIN_NAME: config})
+#
+#
+#@pytest.mark.parametrize("use_discrete", [True, False])
+#@pytest.mark.parametrize("num_visual", [1, 2])
+#def test_visual_ppo(num_visual, use_discrete):
+#    env = SimpleEnvironment(
+#        [BRAIN_NAME],
+#        use_discrete=use_discrete,
+#        num_visual=num_visual,
+#        num_vector=0,
+#        step_size=0.2,
+#    )
+#    new_hyperparams = attr.evolve(PPO_CONFIG.hyperparameters, learning_rate=3.0e-4)
+#    config = attr.evolve(PPO_CONFIG, hyperparameters=new_hyperparams)
+#    _check_environment_trains(env, {BRAIN_NAME: config})
+#
+#
+#@pytest.mark.parametrize("num_visual", [1, 2])
+#@pytest.mark.parametrize("vis_encode_type", ["resnet", "nature_cnn", "match3"])
+#def test_visual_advanced_ppo(vis_encode_type, num_visual):
+#    env = SimpleEnvironment(
+#        [BRAIN_NAME],
+#        use_discrete=True,
+#        num_visual=num_visual,
+#        num_vector=0,
+#        step_size=0.5,
+#        vis_obs_size=(5, 5, 5) if vis_encode_type == "match3" else (36, 36, 3),
+#    )
+#    new_networksettings = attr.evolve(
+#        SAC_CONFIG.network_settings, vis_encode_type=EncoderType(vis_encode_type)
+#    )
+#    new_hyperparams = attr.evolve(PPO_CONFIG.hyperparameters, learning_rate=3.0e-4)
+#    config = attr.evolve(
+#        PPO_CONFIG,
+#        hyperparameters=new_hyperparams,
+#        network_settings=new_networksettings,
+#        max_steps=700,
+#        summary_freq=100,
+#    )
+#    # The number of steps is pretty small for these encoders
+#    _check_environment_trains(env, {BRAIN_NAME: config}, success_threshold=0.5)
+#
+#
+#@pytest.mark.parametrize("use_discrete", [True, False])
+#def test_recurrent_ppo(use_discrete):
+#    env = MemoryEnvironment([BRAIN_NAME], use_discrete=use_discrete)
+#    new_network_settings = attr.evolve(
+#        PPO_CONFIG.network_settings,
+#        memory=NetworkSettings.MemorySettings(memory_size=16),
+#    )
+#    new_hyperparams = attr.evolve(
+#        PPO_CONFIG.hyperparameters, learning_rate=1.0e-3, batch_size=64, buffer_size=128
+#    )
+#    config = attr.evolve(
+#        PPO_CONFIG,
+#        hyperparameters=new_hyperparams,
+#        network_settings=new_network_settings,
+#        max_steps=5000,
+#    )
+#    _check_environment_trains(env, {BRAIN_NAME: config}, success_threshold=0.9)
+#
+#
+#@pytest.mark.parametrize("use_discrete", [True, False])
+#def test_simple_sac(use_discrete):
+#    env = SimpleEnvironment([BRAIN_NAME], use_discrete=use_discrete)
+#    config = attr.evolve(SAC_CONFIG)
+#    _check_environment_trains(env, {BRAIN_NAME: config})
+#
+#
+#@pytest.mark.parametrize("use_discrete", [True, False])
+#def test_2d_sac(use_discrete):
+#    env = SimpleEnvironment(
+#        [BRAIN_NAME], use_discrete=use_discrete, action_size=2, step_size=0.8
+#    )
+#    new_hyperparams = attr.evolve(SAC_CONFIG.hyperparameters, buffer_init_steps=2000)
+#    config = attr.evolve(SAC_CONFIG, hyperparameters=new_hyperparams, max_steps=10000)
+#    _check_environment_trains(env, {BRAIN_NAME: config}, success_threshold=0.8)
+#
+#
+#@pytest.mark.parametrize("use_discrete", [True, False])
+#@pytest.mark.parametrize("num_visual", [1, 2])
+#def test_visual_sac(num_visual, use_discrete):
+#    env = SimpleEnvironment(
+#        [BRAIN_NAME],
+#        use_discrete=use_discrete,
+#        num_visual=num_visual,
+#        num_vector=0,
+#        step_size=0.2,
+#    )
+#    new_hyperparams = attr.evolve(
+#        SAC_CONFIG.hyperparameters, batch_size=16, learning_rate=3e-4
+#    )
+#    config = attr.evolve(SAC_CONFIG, hyperparameters=new_hyperparams)
+#    _check_environment_trains(env, {BRAIN_NAME: config})
+#
+#
+#@pytest.mark.parametrize("num_visual", [1, 2])
+#@pytest.mark.parametrize("vis_encode_type", ["resnet", "nature_cnn", "match3"])
+#def test_visual_advanced_sac(vis_encode_type, num_visual):
+#    env = SimpleEnvironment(
+#        [BRAIN_NAME],
+#        use_discrete=True,
+#        num_visual=num_visual,
+#        num_vector=0,
+#        step_size=0.5,
+#        vis_obs_size=(5, 5, 5) if vis_encode_type == "match3" else (36, 36, 3),
+#    )
+#    new_networksettings = attr.evolve(
+#        SAC_CONFIG.network_settings, vis_encode_type=EncoderType(vis_encode_type)
+#    )
+#    new_hyperparams = attr.evolve(
+#        SAC_CONFIG.hyperparameters,
+#        batch_size=16,
+#        learning_rate=3e-4,
+#        buffer_init_steps=0,
+#    )
+#    config = attr.evolve(
+#        SAC_CONFIG,
+#        hyperparameters=new_hyperparams,
+#        network_settings=new_networksettings,
+#        max_steps=100,
+#    )
+#    # The number of steps is pretty small for these encoders
+#    _check_environment_trains(env, {BRAIN_NAME: config}, success_threshold=0.5)
+#
+#
+#@pytest.mark.parametrize("use_discrete", [True, False])
+#def test_recurrent_sac(use_discrete):
+#    step_size = 0.2 if use_discrete else 0.5
+#    env = MemoryEnvironment(
+#        [BRAIN_NAME], use_discrete=use_discrete, step_size=step_size
+#    )
+#    new_networksettings = attr.evolve(
+#        SAC_CONFIG.network_settings,
+#        memory=NetworkSettings.MemorySettings(memory_size=16, sequence_length=16),
+#    )
+#    new_hyperparams = attr.evolve(
+#        SAC_CONFIG.hyperparameters,
+#        batch_size=128,
+#        learning_rate=1e-3,
+#        buffer_init_steps=1000,
+#        steps_per_update=2,
+#    )
+#    config = attr.evolve(
+#        SAC_CONFIG,
+#        hyperparameters=new_hyperparams,
+#        network_settings=new_networksettings,
+#        max_steps=5000,
+#    )
+#    _check_environment_trains(env, {BRAIN_NAME: config})
+#
+#
+#@pytest.mark.parametrize("use_discrete", [True, False])
+#def test_simple_ghost(use_discrete):
+#    env = SimpleEnvironment(
+#        [BRAIN_NAME + "?team=0", BRAIN_NAME + "?team=1"], use_discrete=use_discrete
+#    )
+#    self_play_settings = SelfPlaySettings(
+#        play_against_latest_model_ratio=1.0, save_steps=2000, swap_steps=2000
+#    )
+#    config = attr.evolve(PPO_CONFIG, self_play=self_play_settings, max_steps=2500)
+#    _check_environment_trains(env, {BRAIN_NAME: config})
+#
+#
+#@pytest.mark.parametrize("use_discrete", [True, False])
+#def test_simple_ghost_fails(use_discrete):
+#    env = SimpleEnvironment(
+#        [BRAIN_NAME + "?team=0", BRAIN_NAME + "?team=1"], use_discrete=use_discrete
+#    )
+#    # This config should fail because the ghosted policy is never swapped with a competent policy.
+#    # Swap occurs after max step is reached.
+#    self_play_settings = SelfPlaySettings(
+#        play_against_latest_model_ratio=1.0, save_steps=2000, swap_steps=4000
+#    )
+#    config = attr.evolve(PPO_CONFIG, self_play=self_play_settings, max_steps=2500)
+#    _check_environment_trains(env, {BRAIN_NAME: config}, success_threshold=None)
+#    processed_rewards = [
+#        default_reward_processor(rewards) for rewards in env.final_rewards.values()
+#    ]
+#    success_threshold = 0.9
+#    assert any(reward > success_threshold for reward in processed_rewards) and any(
+#        reward < success_threshold for reward in processed_rewards
+#    )
+#
+#
+#@pytest.mark.parametrize("use_discrete", [True, False])
+#def test_simple_asymm_ghost(use_discrete):
+#    # Make opponent for asymmetric case
+#    brain_name_opp = BRAIN_NAME + "Opp"
+#    env = SimpleEnvironment(
+#        [BRAIN_NAME + "?team=0", brain_name_opp + "?team=1"], use_discrete=use_discrete
+#    )
+#    self_play_settings = SelfPlaySettings(
+#        play_against_latest_model_ratio=1.0,
+#        save_steps=10000,
+#        swap_steps=10000,
+#        team_change=400,
+#    )
+#    config = attr.evolve(PPO_CONFIG, self_play=self_play_settings, max_steps=4000)
+#    _check_environment_trains(env, {BRAIN_NAME: config, brain_name_opp: config})
+#
+#
+#@pytest.mark.parametrize("use_discrete", [True, False])
+#def test_simple_asymm_ghost_fails(use_discrete):
+#    # Make opponent for asymmetric case
+#    brain_name_opp = BRAIN_NAME + "Opp"
+#    env = SimpleEnvironment(
+#        [BRAIN_NAME + "?team=0", brain_name_opp + "?team=1"], use_discrete=use_discrete
+#    )
+#    # This config should fail because the team that us not learning when both have reached
+#    # max step should be executing the initial, untrained poliy.
+#    self_play_settings = SelfPlaySettings(
+#        play_against_latest_model_ratio=0.0,
+#        save_steps=5000,
+#        swap_steps=5000,
+#        team_change=2000,
+#    )
+#    config = attr.evolve(PPO_CONFIG, self_play=self_play_settings, max_steps=3000)
+#    _check_environment_trains(
+#        env, {BRAIN_NAME: config, brain_name_opp: config}, success_threshold=None
+#    )
+#    processed_rewards = [
+#        default_reward_processor(rewards) for rewards in env.final_rewards.values()
+#    ]
+#    success_threshold = 0.9
+#    assert any(reward > success_threshold for reward in processed_rewards) and any(
+#        reward < success_threshold for reward in processed_rewards
+#    )
+#
+#
+#@pytest.fixture(scope="session")
+#def simple_record(tmpdir_factory):
+#    def record_demo(use_discrete, num_visual=0, num_vector=1):
+#        env = RecordEnvironment(
+#            [BRAIN_NAME],
+#            use_discrete=use_discrete,
+#            num_visual=num_visual,
+#            num_vector=num_vector,
+#            n_demos=100,
+#        )
+#        # If we want to use true demos, we can solve the env in the usual way
+#        # Otherwise, we can just call solve to execute the optimal policy
+#        env.solve()
+#        agent_info_protos = env.demonstration_protos[BRAIN_NAME]
+#        meta_data_proto = DemonstrationMetaProto()
+#        brain_param_proto = BrainParametersProto(
+#            vector_action_size=[2] if use_discrete else [1],
+#            vector_action_descriptions=[""],
+#            vector_action_space_type=discrete if use_discrete else continuous,
+#            brain_name=BRAIN_NAME,
+#            is_training=True,
+#        )
+#        action_type = "Discrete" if use_discrete else "Continuous"
+#        demo_path_name = "1DTest" + action_type + ".demo"
+#        demo_path = str(tmpdir_factory.mktemp("tmp_demo").join(demo_path_name))
+#        write_demo(demo_path, meta_data_proto, brain_param_proto, agent_info_protos)
+#        return demo_path
+#
+#    return record_demo
+#
+#
+#@pytest.mark.parametrize("use_discrete", [True, False])
+#@pytest.mark.parametrize("trainer_config", [PPO_CONFIG, SAC_CONFIG])
+#def test_gail(simple_record, use_discrete, trainer_config):
+#    demo_path = simple_record(use_discrete)
+#    env = SimpleEnvironment([BRAIN_NAME], use_discrete=use_discrete, step_size=0.2)
+#    bc_settings = BehavioralCloningSettings(demo_path=demo_path, steps=1000)
+#    reward_signals = {
+#        RewardSignalType.GAIL: GAILSettings(encoding_size=32, demo_path=demo_path)
+#    }
+#    config = attr.evolve(
+#        trainer_config,
+#        reward_signals=reward_signals,
+#        behavioral_cloning=bc_settings,
+#        max_steps=500,
+#    )
+#    _check_environment_trains(env, {BRAIN_NAME: config}, success_threshold=0.9)
+#
+#
+#@pytest.mark.parametrize("use_discrete", [True, False])
+#def test_gail_visual_ppo(simple_record, use_discrete):
+#    demo_path = simple_record(use_discrete, num_visual=1, num_vector=0)
+#    env = SimpleEnvironment(
+#        [BRAIN_NAME],
+#        num_visual=1,
+#        num_vector=0,
+#        use_discrete=use_discrete,
+#        step_size=0.2,
+#    )
+#    bc_settings = BehavioralCloningSettings(demo_path=demo_path, steps=1500)
+#    reward_signals = {
+#        RewardSignalType.GAIL: GAILSettings(encoding_size=32, demo_path=demo_path)
+#    }
+#    hyperparams = attr.evolve(PPO_CONFIG.hyperparameters, learning_rate=3e-4)
+#    config = attr.evolve(
+#        PPO_CONFIG,
+#        reward_signals=reward_signals,
+#        hyperparameters=hyperparams,
+#        behavioral_cloning=bc_settings,
+#        max_steps=1000,
+#    )
+#    _check_environment_trains(env, {BRAIN_NAME: config}, success_threshold=0.9)
+#
+#
+#@pytest.mark.parametrize("use_discrete", [True, False])
+#def test_gail_visual_sac(simple_record, use_discrete):
+#    demo_path = simple_record(use_discrete, num_visual=1, num_vector=0)
+#    env = SimpleEnvironment(
+#        [BRAIN_NAME],
+#        num_visual=1,
+#        num_vector=0,
+#        use_discrete=use_discrete,
+#        step_size=0.2,
+#    )
+#    bc_settings = BehavioralCloningSettings(demo_path=demo_path, steps=1000)
+#    reward_signals = {
+#        RewardSignalType.GAIL: GAILSettings(encoding_size=32, demo_path=demo_path)
+#    }
+#    hyperparams = attr.evolve(
+#        SAC_CONFIG.hyperparameters, learning_rate=3e-4, batch_size=16
+#    )
+#    config = attr.evolve(
+#        SAC_CONFIG,
+#        reward_signals=reward_signals,
+#        hyperparameters=hyperparams,
+#        behavioral_cloning=bc_settings,
+#        max_steps=500,
+#    )
+#    _check_environment_trains(env, {BRAIN_NAME: config}, success_threshold=0.9)

ml-agents/mlagents/trainers/torch/distributions.py

         self,
         hidden_size: int,
         continuous_act_size: int,
-        discrete_act_size: int,
+        discrete_act_size: List[int],
+
+        self.encoding_size = hidden_size
         self.continuous_distributions: List[GaussianDistribution] = []
         self.discrete_distributions: List[MultiCategoricalDistribution] = []
         if continuous_act_size > 0:
                     tanh_squash=tanh_squash,
                 )
             )
-        if discrete_act_size > 0:
+        if len(discrete_act_size) > 0:
             self.discrete_distributions.append(
                 MultiCategoricalDistribution(self.encoding_size, discrete_act_size)
             )

ml-agents/mlagents/trainers/torch/networks.py

 
         self.output_distributions = OutputDistributions(
             self.encoding_size,
-            continuous_act_size[0],
+            continuous_act_size,
             discrete_act_size,
             conditional_sigma=conditional_sigma,
             tanh_squash=tanh_squash,
         tanh_squash: bool = False,
     ):
         super().__init__(
-            self,
-            act_type,
-            act_size,
+            continuous_act_size,
+            discrete_act_size,
             conditional_sigma,
             tanh_squash,
         )
         tanh_squash: bool = False,
     ):
         super().__init__(
-            self,
-            act_type,
-            act_size,
+            continuous_act_size,
+            discrete_act_size,
+        print("CREATED", self.memory_size)
 
     @property
     def memory_size(self) -> int:
 #        self.critic.network_body.update_normalization(vector_obs)
 #
 #
-# class GlobalSteps(nn.Module):
-#    def __init__(self):
-#        super().__init__()
-#        self.__global_step = nn.Parameter(torch.Tensor([0]), requires_grad=False)
-#
-#    @property
-#    def current_step(self):
-#        return int(self.__global_step.item())
-#
-#    @current_step.setter
-#    def current_step(self, value):
-#        self.__global_step[:] = value
-#
-#    def increment(self, value):
-#        self.__global_step += value
-#
-#
-# class LearningRate(nn.Module):
-#    def __init__(self, lr):
-#        # Todo: add learning rate decay
-#        super().__init__()
-#        self.learning_rate = torch.Tensor([lr])
+class GlobalSteps(nn.Module):
+   def __init__(self):
+       super().__init__()
+       self.__global_step = nn.Parameter(torch.Tensor([0]), requires_grad=False)
+
+   @property
+   def current_step(self):
+       return int(self.__global_step.item())
+
+   @current_step.setter
+   def current_step(self, value):
+       self.__global_step[:] = value
+
+   def increment(self, value):
+       self.__global_step += value
+
+
+class LearningRate(nn.Module):
+   def __init__(self, lr):
+       # Todo: add learning rate decay
+       super().__init__()
+       self.learning_rate = torch.Tensor([lr])