Fixed dataset from_array

- added missing policy state in the new dataset
MushroomRL · Nov 26, 2024 · 5395905 · 5395905
1 parent 5c3834d
commit 5395905
Showing 1 changed file with 6 additions and 3 deletions.
diff --git a/mushroom_rl/core/dataset.py b/mushroom_rl/core/dataset.py
@@ -212,11 +212,14 @@ def from_array(cls, states, actions, rewards, next_states, absorbings, lasts,
 
         dataset._array_backend = ArrayBackend.get_array_backend(backend)
         if backend == 'numpy':
-            dataset._data = NumpyDataset.from_array(states, actions, rewards, next_states, absorbings, lasts)
+            dataset._data = NumpyDataset.from_array(states, actions, rewards, next_states, absorbings, lasts,
+                                                    policy_state, policy_next_state)
         elif backend == 'torch':
-            dataset._data = TorchDataset.from_array(states, actions, rewards, next_states, absorbings, lasts)
+            dataset._data = TorchDataset.from_array(states, actions, rewards, next_states, absorbings, lasts,
+                                                    policy_state, policy_next_state)
         else:
-            dataset._data = ListDataset.from_array(states, actions, rewards, next_states, absorbings, lasts)
+            dataset._data = ListDataset.from_array(states, actions, rewards, next_states, absorbings, lasts,
+                                                   policy_state, policy_next_state)
 
         state_shape = states.shape[1:]
         action_shape = actions.shape[1:]