Fixed minor bugs in tests

MushroomRL · Feb 5, 2025 · 24fe268 · 24fe268
1 parent 2bc7612
commit 24fe268
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 3 deletions.
diff --git a/tests/rl_utils/test_value_functions.py b/tests/rl_utils/test_value_functions.py
@@ -1,4 +1,5 @@
 import torch
+import numpy as np
 from mushroom_rl.policy import DeterministicPolicy
 from mushroom_rl.environments.segway import Segway
 from mushroom_rl.core import Core, Agent
@@ -49,7 +50,8 @@ def gae(V, s, ss, r, absorbing, last, gamma, lam):
 
 def _value_functions_tester(test_fun, correct_fun, *args):
     mdp = Segway()
-    V = Regressor(TorchApproximator, input_shape=mdp.info.observation_space.shape, output_shape=(1,), network=Net, loss=torch.nn.MSELoss(), optimizer={'class': torch.optim.Adam, 'params': {'lr': 0.001}})
+    V = Regressor(TorchApproximator, input_shape=mdp.info.observation_space.shape, output_shape=(1,),
+                  network=Net, loss=torch.nn.MSELoss(), optimizer={'class': torch.optim.Adam, 'params': {'lr': 0.001}})
 
     state, action, reward, next_state, absorbing, last = _get_episodes(mdp, 10)
 
@@ -68,7 +70,7 @@ def _value_functions_tester(test_fun, correct_fun, *args):
     assert torch.allclose(adv, correct_adv)
 
 def _get_episodes(mdp, n_episodes=100):
-    mu = torch.tensor([6.31154476, 3.32346271, 0.49648221]).unsqueeze(0)
+    mu = np.array([6.31154476, 3.32346271, 0.49648221])
 
     approximator = Regressor(LinearApproximator,
                             input_shape=mdp.info.observation_space.shape,

diff --git a/tests/utils/test_episodes.py b/tests/utils/test_episodes.py
@@ -44,7 +44,7 @@ def test_numpy_split():
     assert np.allclose(last, un_last)
 
 def get_episodes(mdp, n_episodes=100):
-    mu = torch.tensor([6.31154476, 3.32346271, 0.49648221]).unsqueeze(0)
+    mu = np.array([6.31154476, 3.32346271, 0.49648221])
 
     approximator = Regressor(LinearApproximator,
                             input_shape=mdp.info.observation_space.shape,