Revert "Add hybrid actions to SAC"

This reverts commit a759b36a51df4f8f1fd296f9f148269f0f026e42.

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

         self, decision_requests: DecisionSteps
     ) -> Tuple[SplitObservations, np.ndarray]:
         vec_vis_obs = SplitObservations.from_observations(decision_requests.obs)
-        # mask = None
+        #mask = None
         mask = torch.ones([len(decision_requests), np.sum(self.discrete_act_size)])
         if decision_requests.action_mask is not None:
             mask = torch.as_tensor(
         :param all_log_probs: Returns (for discrete actions) a tensor of log probs, one for each action.
         """
         actions, log_probs, entropies, value_heads, memories = self.actor_critic.get_action_stats_and_value(
-            vec_obs, vis_obs, masks, memories, seq_len, all_log_probs
+            vec_obs, vis_obs, masks, memories, seq_len
         )
         return (
             actions,

ml-agents/mlagents/trainers/sac/optimizer_torch.py

-from mlagents_envs.communicator_objects.space_type_pb2 import continuous
 import numpy as np
 from typing import Dict, List, Mapping, cast, Tuple, Optional
 from mlagents.torch_utils import torch, nn, default_device
             stream_names: List[str],
             observation_shapes: List[Tuple[int, ...]],
             network_settings: NetworkSettings,
-            continuous_act_size: List[int],
-            discrete_act_size: List[int],
+            act_type: ActionType,
+            act_size: List[int],
-            num_value_outs = max(sum(discrete_act_size), 1)
-            num_action_ins = int(continuous_act_size)
-
+            if act_type == ActionType.CONTINUOUS:
+                num_value_outs = 1
+                num_action_ins = sum(act_size)
+            else:
+                num_value_outs = sum(act_size)
+                num_action_ins = 0
             self.q1_network = ValueNetwork(
                 stream_names,
                 observation_shapes,
         self.init_entcoef = hyperparameters.init_entcoef
 
         self.policy = policy
+        self.act_size = policy.act_size
         policy_network_settings = policy.network_settings
 
         self.tau = hyperparameters.tau
             self.stream_names,
             self.policy.behavior_spec.observation_shapes,
             policy_network_settings,
-            self.policy.continuous_act_size,
-            self.policy.discrete_act_size,
+            self.policy.behavior_spec.action_type,
+            self.act_size,
         )
 
         self.target_network = ValueNetwork(
         )
         self.soft_update(self.policy.actor_critic.critic, self.target_network, 1.0)
 
-        _total_act_size = 0
-        if self.policy.continuous_act_size > 0:
-            _total_act_size += 1
-        _total_act_size += len(self.policy.discrete_act_size)
-
-            torch.log(torch.as_tensor([self.init_entcoef] * _total_act_size)),
+            torch.log(torch.as_tensor([self.init_entcoef] * len(self.act_size))),
-        self.target_entropy = []
-        if self.policy.continuous_act_size > 0:
-            self.target_entropy.append(
-                torch.as_tensor(
-                    -1
-                    * self.continuous_target_entropy_scale
-                    * np.prod(self.policy.continuous_act_size).astype(np.float32)
-                )
+        if self.policy.use_continuous_act:
+            self.target_entropy = torch.as_tensor(
+                -1
+                * self.continuous_target_entropy_scale
+                * np.prod(self.act_size[0]).astype(np.float32)
-            self.target_entropy += [
+            self.target_entropy = [
-                for i in self.policy.discrete_act_size
+                for i in self.act_size
-            self.policy.actor_critic.action_model.parameters()
+            self.policy.actor_critic.distribution.parameters()
         )
         value_params = list(self.value_network.parameters()) + list(
             self.policy.actor_critic.critic.parameters()
         q1p_out: Dict[str, torch.Tensor],
         q2p_out: Dict[str, torch.Tensor],
         loss_masks: torch.Tensor,
+        discrete: bool,
-            if len(self.policy.discrete_act_size) <= 0:
+            if not discrete:
-                disc_action_probs = log_probs[:, self.policy.continuous_act_size:].exp()
+                action_probs = log_probs.exp()
-                    q1p_out[name] * disc_action_probs, self.policy.discrete_act_size
+                    q1p_out[name] * action_probs, self.act_size
-                    q2p_out[name] * disc_action_probs, self.policy.discrete_act_size
+                    q2p_out[name] * action_probs, self.act_size
                 )
                 _q1p_mean = torch.mean(
                     torch.stack(
                 min_policy_qs[name] = torch.min(_q1p_mean, _q2p_mean)
 
         value_losses = []
-        if len(self.policy.discrete_act_size) <= 0:
+        if not discrete:
             for name in values.keys():
                 with torch.no_grad():
                     v_backup = min_policy_qs[name] - torch.sum(
                 value_losses.append(value_loss)
         else:
             branched_per_action_ent = ModelUtils.break_into_branches(
-                log_probs * log_probs.exp(), self.policy.discrete_act_size
+                log_probs * log_probs.exp(), self.act_size
             )
             # We have to do entropy bonus per action branch
             branched_ent_bonus = torch.stack(
         log_probs: torch.Tensor,
         q1p_outs: Dict[str, torch.Tensor],
         loss_masks: torch.Tensor,
+        discrete: bool,
-        if len(self.policy.discrete_act_size) <= 0:
+        if not discrete:
-            disc_log_probs = log_probs[:, self.policy.continuous_act_size:]
-            disc_action_probs = disc_log_probs.exp()
+            action_probs = log_probs.exp()
-                disc_log_probs * disc_action_probs, self.policy.discrete_act_size
+                log_probs * action_probs, self.act_size
-                mean_q1 * disc_action_probs, self.policy.discrete_act_size
+                mean_q1 * action_probs, self.act_size
             )
             branched_policy_loss = torch.stack(
                 [
         return policy_loss
 
     def sac_entropy_loss(
-        self, log_probs: torch.Tensor, loss_masks: torch.Tensor
+        self, log_probs: torch.Tensor, loss_masks: torch.Tensor, discrete: bool
-        if len(self.policy.discrete_act_size) <= 0:
+        if not discrete:
             with torch.no_grad():
                 target_current_diff = torch.sum(log_probs + self.target_entropy, dim=1)
             entropy_loss = -torch.mean(
             with torch.no_grad():
                 branched_per_action_ent = ModelUtils.break_into_branches(
-                    log_probs * log_probs.exp(), self.policy.discrete_act_size
+                    log_probs * log_probs.exp(), self.act_size
                 )
                 target_current_diff_branched = torch.stack(
                     [
         self, q_output: Dict[str, torch.Tensor], discrete_actions: torch.Tensor
     ) -> Dict[str, torch.Tensor]:
         condensed_q_output = {}
-        onehot_actions = ModelUtils.actions_to_onehot(
-            discrete_actions, self.policy.discrete_act_size
-        )
+        onehot_actions = ModelUtils.actions_to_onehot(discrete_actions, self.act_size)
-            branched_q = ModelUtils.break_into_branches(
-                item, self.policy.discrete_act_size
-            )
+            branched_q = ModelUtils.break_into_branches(item, self.act_size)
             only_action_qs = torch.stack(
                 [
                     torch.sum(_act * _q, dim=1, keepdim=True)
         vec_obs = [ModelUtils.list_to_tensor(batch["vector_obs"])]
         next_vec_obs = [ModelUtils.list_to_tensor(batch["next_vector_in"])]
         act_masks = ModelUtils.list_to_tensor(batch["action_mask"])
-
-        actions = ModelUtils.list_to_tensor(batch["actions"]).unsqueeze(-1)
+        if self.policy.use_continuous_act:
+            actions = ModelUtils.list_to_tensor(batch["actions"]).unsqueeze(-1)
+        else:
+            actions = ModelUtils.list_to_tensor(batch["actions"], dtype=torch.long)
 
         memories_list = [
             ModelUtils.list_to_tensor(batch["memory"][i])
             masks=act_masks,
             memories=memories,
             seq_len=self.policy.sequence_length,
-            all_log_probs=True,
+            all_log_probs=not self.policy.use_continuous_act,
-
-        cont_sampled_actions = sampled_actions[:, : self.policy.continuous_act_size]
-        # disc_sampled_actions = sampled_actions[:, self.policy.continuous_act_size :]
-
-        cont_actions = actions[:, : self.policy.continuous_act_size]
-        disc_actions = actions[:, self.policy.continuous_act_size :]
-        q1p_out, q2p_out = self.value_network(
-            vec_obs,
-            vis_obs,
-            cont_sampled_actions,
-            memories=q_memories,
-            sequence_length=self.policy.sequence_length,
-        )
-        q1_out, q2_out = self.value_network(
-            vec_obs,
-            vis_obs,
-            cont_actions.squeeze(-1),
-            memories=q_memories,
-            sequence_length=self.policy.sequence_length,
-        )
-        if self.policy.discrete_act_size:
-            q1_stream = self._condense_q_streams(q1_out, disc_actions.squeeze(-1))
-            q2_stream = self._condense_q_streams(q2_out, disc_actions.squeeze(-1))
-        else:
+        if self.policy.use_continuous_act:
+            squeezed_actions = actions.squeeze(-1)
+            q1p_out, q2p_out = self.value_network(
+                vec_obs,
+                vis_obs,
+                sampled_actions,
+                memories=q_memories,
+                sequence_length=self.policy.sequence_length,
+            )
+            q1_out, q2_out = self.value_network(
+                vec_obs,
+                vis_obs,
+                squeezed_actions,
+                memories=q_memories,
+                sequence_length=self.policy.sequence_length,
+            )
+        else:
+            with torch.no_grad():
+                q1p_out, q2p_out = self.value_network(
+                    vec_obs,
+                    vis_obs,
+                    memories=q_memories,
+                    sequence_length=self.policy.sequence_length,
+                )
+            q1_out, q2_out = self.value_network(
+                vec_obs,
+                vis_obs,
+                memories=q_memories,
+                sequence_length=self.policy.sequence_length,
+            )
+            q1_stream = self._condense_q_streams(q1_out, actions)
+            q2_stream = self._condense_q_streams(q2_out, actions)
 
         with torch.no_grad():
             target_values, _ = self.target_network(
                 sequence_length=self.policy.sequence_length,
             )
         masks = ModelUtils.list_to_tensor(batch["masks"], dtype=torch.bool)
+        use_discrete = not self.policy.use_continuous_act
         dones = ModelUtils.list_to_tensor(batch["done"])
 
         q1_loss, q2_loss = self.sac_q_loss(
-            log_probs, sampled_values, q1p_out, q2p_out, masks
+            log_probs, sampled_values, q1p_out, q2p_out, masks, use_discrete
-        policy_loss = self.sac_policy_loss(log_probs, q1p_out, masks)
-        entropy_loss = self.sac_entropy_loss(log_probs, masks)
+        policy_loss = self.sac_policy_loss(log_probs, q1p_out, masks, use_discrete)
+        entropy_loss = self.sac_entropy_loss(log_probs, masks, use_discrete)
 
         total_value_loss = q1_loss + q2_loss + value_loss
 

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

     config = attr.evolve(PPO_CONFIG, hyperparameters=new_hyperparams, max_steps=100000)
     _check_environment_trains(env, {BRAIN_NAME: config}, success_threshold=5.0)
 
-def test_2dhybrid_sac():
-    env = HybridEnvironment([BRAIN_NAME], continuous_action_size=1, discrete_action_size=2, step_size=0.8)
-    new_hyperparams = attr.evolve(
-        SAC_CONFIG.hyperparameters, buffer_size=50000, batch_size=128
-    )
-    config = attr.evolve(SAC_CONFIG, hyperparameters=new_hyperparams, max_steps=100000)
-    _check_environment_trains(env, {BRAIN_NAME: config}, success_threshold=5.0)
-
 #
 #@pytest.mark.parametrize("use_discrete", [True, False])
 #def test_2d_ppo(use_discrete):

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

 import numpy as np
 import math
 from mlagents.trainers.torch.layers import linear_layer, Initialization
-from mlagents.trainers.torch.distributions import (
-    DistInstance,
-    DiscreteDistInstance,
-    GaussianDistribution,
-    MultiCategoricalDistribution,
-)
+from mlagents.trainers.torch.distributions import DistInstance, DiscreteDistInstance, GaussianDistribution, MultiCategoricalDistribution
 
 from mlagents.trainers.torch.utils import ModelUtils
 
+
+
-    # @abc.abstractmethod
-    # def entropy(self, action_list: np.ndarray) -> torch.Tensor:
+    #@abc.abstractmethod
+    #def entropy(self, action_list: np.ndarray) -> torch.Tensor:
-    # @abc.abstractmethod
-    # def log_probs(self, action_list: np.ndarray) -> torch.Tensor:
+    #@abc.abstractmethod
+    #def log_probs(self, action_list: np.ndarray) -> torch.Tensor:
     #    pass
 
     def _sample_action(self, dists: List[DistInstance]) -> List[torch.Tensor]:
     @abc.abstractmethod
     def forward(self, inputs: torch.Tensor, masks: torch.Tensor):
         pass
-
 
 class HybridActionModel(ActionModel):
     def __init__(
         self.encoding_size = hidden_size
         self.continuous_act_size = continuous_act_size
         self.discrete_act_size = discrete_act_size
-        self.continuous_distribution = None  #: List[GaussianDistribution] = []
-        self.discrete_distribution = None  #: List[MultiCategoricalDistribution] = []
+        self.continuous_distribution = None #: List[GaussianDistribution] = []
+        self.discrete_distribution = None #: List[MultiCategoricalDistribution] = []
-                self.encoding_size,
-                continuous_act_size,
-                conditional_sigma=conditional_sigma,
-                tanh_squash=tanh_squash,
-            )
+                    self.encoding_size,
+                    continuous_act_size,
+                    conditional_sigma=conditional_sigma,
+                    tanh_squash=tanh_squash,
+                )
-            self.discrete_distribution = MultiCategoricalDistribution(
-                self.encoding_size, discrete_act_size
-            )
+            self.discrete_distribution = MultiCategoricalDistribution(self.encoding_size, discrete_act_size)
+            
-    def evaluate(
-        self, inputs: torch.Tensor, masks: torch.Tensor, actions: torch.Tensor
-    ) -> Tuple[torch.Tensor, torch.Tensor]:
+    def evaluate(self, inputs: torch.Tensor, masks: torch.Tensor, actions: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
-        continuous_actions, discrete_actions = torch.split(
-            actions, [self.continuous_act_size, len(self.discrete_act_size)], dim=1
-        )
+        continuous_actions, discrete_actions = torch.split(actions, [self.continuous_act_size, len(self.discrete_act_size)], dim=1) 
-        continuous_action_list = [
-            continuous_actions[..., i] for i in range(continuous_actions.shape[-1])
-        ]
-        (
-            continuous_log_probs,
-            continuous_entropies,
-            _,
-        ) = ModelUtils.get_probs_and_entropy(continuous_action_list, continuous_dists)
+
+        continuous_action_list = [continuous_actions[..., i] for i in range(continuous_actions.shape[-1])]
+        continuous_log_probs, continuous_entropies, _ = ModelUtils.get_probs_and_entropy(continuous_action_list, continuous_dists)
-        discrete_action_list = [
-            discrete_actions[:, i] for i in range(len(self.discrete_act_size))
-        ]
-        discrete_log_probs, discrete_entropies, _ = ModelUtils.get_probs_and_entropy(
-            discrete_action_list, discrete_dists
-        )
+        discrete_action_list = [discrete_actions[:, i] for i in range(len(self.discrete_act_size))]
+        discrete_log_probs, discrete_entropies, _ = ModelUtils.get_probs_and_entropy(discrete_action_list, discrete_dists)
-        log_probs = torch.cat([continuous_log_probs, discrete_log_probs], dim=1)
-        entropies = torch.cat(
-            [continuous_entropies, torch.mean(discrete_entropies, dim=0).unsqueeze(0)],
-            dim=1,
-        )
+        log_probs = torch.cat([continuous_log_probs, discrete_log_probs], dim=1) 
+        entropies = torch.cat([continuous_entropies, torch.mean(discrete_entropies, dim=0).unsqueeze(0)], dim=1) 
         return log_probs, entropies
 
     def get_action_out(self, inputs: torch.Tensor, masks: torch.Tensor) -> torch.Tensor:
 
-    def _get_dists(
-        self, inputs: torch.Tensor, masks: torch.Tensor
-    ) -> Tuple[List[DistInstance], List[DiscreteDistInstance]]:
-        # continuous_distributions: List[DistInstance] = []
-        # discrete_distributions: List[DiscreteDistInstance] = []
-        continuous_dist_instances = self.continuous_distribution(
-            inputs
-        )  # for continuous_dist in self.continuous_distributions]
-        discrete_dist_instances = self.discrete_distribution(
-            inputs, masks
-        )  # for discrete_dist in self.discrete_distributions]
-        # for continuous_dist in self.continuous_distributions:
+    def _get_dists(self, inputs: torch.Tensor, masks: torch.Tensor) -> Tuple[List[DistInstance], List[DiscreteDistInstance]]:
+        #continuous_distributions: List[DistInstance] = []
+        #discrete_distributions: List[DiscreteDistInstance] = []
+        continuous_dist_instances = self.continuous_distribution(inputs)# for continuous_dist in self.continuous_distributions]
+        discrete_dist_instances = self.discrete_distribution(inputs, masks)# for discrete_dist in self.discrete_distributions]
+        #for continuous_dist in self.continuous_distributions:
-        # for discrete_dist in self.discrete_distributions:
+        #for discrete_dist in self.discrete_distributions:
-    def _get_discrete_continuous_outputs(
-        self, inputs: torch.Tensor, masks: torch.Tensor
-    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    def forward(self, inputs: torch.Tensor, masks: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
-        (
-            continuous_entropies,
-            continuous_log_probs,
-            continuous_all_probs,
-        ) = ModelUtils.get_probs_and_entropy(continuous_action_list, continuous_dists)
+        continuous_entropies, continuous_log_probs, continuous_all_probs = ModelUtils.get_probs_and_entropy(
+            continuous_action_list, continuous_dists
+        )
-        (
-            discrete_entropies,
-            discrete_log_probs,
-            discrete_all_probs,
-        ) = ModelUtils.get_probs_and_entropy(discrete_action_list, discrete_dists)
+        discrete_entropies, discrete_log_probs, discrete_all_probs = ModelUtils.get_probs_and_entropy(
+            discrete_action_list, discrete_dists
+        )
-        return (
-            discrete_actions,
-            discrete_entropies,
-            discrete_log_probs,
-            discrete_all_probs,
-            continuous_actions,
-            continuous_entropies,
-            continuous_log_probs,
-            continuous_all_probs,
-        )
-
-    def forward_all_disc_probs(
-        self, inputs: torch.Tensor, masks: torch.Tensor
-    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
-        (
-            discrete_actions,
-            discrete_entropies,
-            _,
-            discrete_all_probs,
-            continuous_actions,
-            continuous_entropies,
-            continuous_log_probs,
-            _,
-        ) = self._get_discrete_continuous_outputs(inputs, masks)
-
-        action = torch.cat(
-            [continuous_actions, discrete_actions.type(torch.float)], dim=1
-        )
-        log_probs = torch.cat([continuous_log_probs, discrete_all_probs], dim=1)
-        entropies = torch.cat([continuous_entropies, discrete_entropies], dim=1)
-        return (action, log_probs, entropies)
-
-    def forward(
-        self, inputs: torch.Tensor, masks: torch.Tensor
-    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
-        (
-            discrete_actions,
-            discrete_entropies,
-            discrete_log_probs,
-            _,
-            continuous_actions,
-            continuous_entropies,
-            continuous_log_probs,
-            _,
-        ) = self._get_discrete_continuous_outputs(inputs, masks)
-
-        action = torch.cat(
-            [continuous_actions, discrete_actions.type(torch.float)], dim=1
-        )
-        log_probs = torch.cat([continuous_log_probs, discrete_log_probs], dim=1)
-        entropies = torch.cat([continuous_entropies, discrete_entropies], dim=1)
+        action = torch.cat([continuous_actions, discrete_actions.type(torch.float)], dim=1) 
+        log_probs = torch.cat([continuous_log_probs, discrete_log_probs], dim=1) 
+        entropies = torch.cat([continuous_entropies, discrete_entropies], dim=1) 
         return (action, log_probs, entropies)

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

         masks: Optional[torch.Tensor] = None,
         memories: Optional[torch.Tensor] = None,
         sequence_length: int = 1,
-        all_discrete_probs: bool = False,
-        if all_discrete_probs:
-            action, log_probs, entropies = self.action_model.forward_all_disc_probs(encoding, masks)
-        else:
-            action, log_probs, entropies = self.action_model(encoding, masks)
+        action, log_probs, entropies = self.action_model(encoding, masks)
         value_outputs = self.value_heads(encoding)
         return action, log_probs, entropies, value_outputs, memories
 
         masks: Optional[torch.Tensor] = None,
         memories: Optional[torch.Tensor] = None,
         sequence_length: int = 1,
-        all_discrete_probs: bool = False,
     ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, Dict[str, torch.Tensor], torch.Tensor]:
         if self.use_lstm:
             # Use only the back half of memories for critic and actor
         encoding, memories = self.network_body(
             vec_inputs, vis_inputs, memories=memories, sequence_length=sequence_length
         )
-        if all_discrete_probs:
-            action, log_probs, entropies = self.action_model.forward_all_disc_probs(encoding, masks)
-        else:
-            action, log_probs, entropies = self.action_model(encoding, masks)
+        action, log_probs, entropies = self.action_model(encoding, masks)
         value_outputs, critic_mem_outs = self.critic(
         vec_inputs, vis_inputs, memories=critic_mem, sequence_length=sequence_length
         )