[refactor] Reduce dimension of continuous ent coef (#4725)

Ervin T · web-flow · commit a2c419de24c5 · 2020-12-08T23:31:52.000-05:00
diff --git a/ml-agents/mlagents/trainers/sac/optimizer_torch.py b/ml-agents/mlagents/trainers/sac/optimizer_torch.py
@@ -159,10 +159,7 @@ def __init__(self, policy: TorchPolicy, trainer_params: TrainerSettings):
             requires_grad=True,
         )
         _cont_log_ent_coef = torch.nn.Parameter(
-            torch.log(
-                torch.as_tensor([self.init_entcoef] * self._action_spec.continuous_size)
-            ),
-            requires_grad=True,
+            torch.log(torch.as_tensor([self.init_entcoef])), requires_grad=True
         )
         self._log_ent_coef = TorchSACOptimizer.LogEntCoef(
             discrete=_disc_log_ent_coef, continuous=_cont_log_ent_coef
@@ -426,7 +423,7 @@ def sac_entropy_loss(
                 )
             # We update all the _cont_ent_coef as one block
             entropy_loss += -1 * ModelUtils.masked_mean(
-                torch.mean(_cont_ent_coef) * target_current_diff, loss_masks
+                _cont_ent_coef * target_current_diff, loss_masks
             )
 
         return entropy_loss

Original file line number	Diff line number	Diff line change
`@@ -159,10 +159,7 @@ def __init__(self, policy: TorchPolicy, trainer_params: TrainerSettings):`
`159`	`159`	`requires_grad=True,`
`160`	`160`	`)`
`161`	`161`	`_cont_log_ent_coef = torch.nn.Parameter(`
`162`		`- torch.log(`
`163`		`- torch.as_tensor([self.init_entcoef] * self._action_spec.continuous_size)`
`164`		`- ),`
`165`		`- requires_grad=True,`
	`162`	`+ torch.log(torch.as_tensor([self.init_entcoef])), requires_grad=True`
`166`	`163`	`)`
`167`	`164`	`self._log_ent_coef = TorchSACOptimizer.LogEntCoef(`
`168`	`165`	`discrete=_disc_log_ent_coef, continuous=_cont_log_ent_coef`
`@@ -426,7 +423,7 @@ def sac_entropy_loss(`
`426`	`423`	`)`
`427`	`424`	`# We update all the _cont_ent_coef as one block`
`428`	`425`	`entropy_loss += -1 * ModelUtils.masked_mean(`
`429`		`- torch.mean(_cont_ent_coef) * target_current_diff, loss_masks`
	`426`	`+ _cont_ent_coef * target_current_diff, loss_masks`
`430`	`427`	`)`
`431`	`428`
`432`	`429`	`return entropy_loss`