Use correct dimensions of gradient

ml-agents/mlagents/trainers/torch/components/reward_providers/gail_reward_provider.py

 
     def evaluate(self, mini_batch: AgentBuffer) -> np.ndarray:
         with torch.no_grad():
+            self._discriminator_network.update_normalization(mini_batch)
             estimates, _ = self._discriminator_network.compute_estimate(
                 mini_batch, use_vail_noise=False
             )
         self._settings = settings
 
         encoder_settings = NetworkSettings(
-            normalize=False,
+            normalize=True,
             hidden_units=settings.encoding_size,
             num_layers=2,
             vis_encode_type=EncoderType.SIMPLE,
         self._estimator = torch.nn.Sequential(
             linear_layer(estimator_input_size, 1), torch.nn.Sigmoid()
         )
+
+    def update_normalization(self, mini_batch: AgentBuffer) -> None:
+        """
+        Updates the normalization of this Discriminator's encoder.
+        """
+        vec_inputs, _ = self.get_state_inputs(mini_batch)
+        self.encoder.update_normalization(vec_inputs)
 
     def get_action_input(self, mini_batch: AgentBuffer) -> torch.Tensor:
         """
             use_vail_noise = True
             z_mu = self._z_mu_layer(hidden)
             hidden = torch.normal(z_mu, self._z_sigma * use_vail_noise)
-        estimate = self._estimator(hidden).squeeze(1).sum()
-        gradient = torch.autograd.grad(estimate, encoder_input, create_graph=True)[0]
+        estimate = self._estimator(hidden).squeeze(1)
+        gradient = torch.autograd.grad(
+            estimate,
+            encoder_input,
+            grad_outputs=torch.ones(estimate.shape),
+            create_graph=True,
+        )[0]
-        safe_norm = (torch.sum(gradient ** 2, dim=1) + self.EPSILON).sqrt()
+        safe_norm = (torch.sum(torch.pow(gradient, 2), dim=1) + self.EPSILON).sqrt()
         gradient_mag = torch.mean((safe_norm - 1) ** 2)
         return gradient_mag