Minor fixes and improvements in core._impl

MushroomRL · Dec 9, 2023 · d5f9e30 · d5f9e30
1 parent d9da4b5
commit d5f9e30
Show file tree

Hide file tree

Showing 4 changed files with 17 additions and 0 deletions.
diff --git a/mushroom_rl/core/_impl/core_logic.py b/mushroom_rl/core/_impl/core_logic.py
@@ -96,3 +96,11 @@ def _move_episodes_condition(self):
 
     def _fit_episodes_condition(self):
         return self._current_episodes_counter >= self._n_episodes_per_fit
+
+    @property
+    def n_steps_per_fit(self):
+        return self._n_steps_per_fit
+
+    @property
+    def n_episodes_per_fit(self):
+        return self._n_episodes_per_fit
diff --git a/mushroom_rl/core/_impl/list_dataset.py b/mushroom_rl/core/_impl/list_dataset.py
@@ -62,6 +62,9 @@ def get_view(self, index):
         else:
             view._dataset = [self._dataset[i] for i in index]
 
+        if self._mask is not None:
+            view._mask = self._mask[index, ...]
+
         return view
 
     def __getitem__(self, index):

diff --git a/mushroom_rl/core/_impl/numpy_dataset.py b/mushroom_rl/core/_impl/numpy_dataset.py
@@ -148,6 +148,9 @@ def get_view(self, index):
             view._policy_states = self._policy_states[index, ...]
             view._policy_next_states = self._policy_next_states[index, ...]
 
+        if self._mask is not None:
+            view._mask = self._mask[index, ...]
+
         return view
 
     def __getitem__(self, index):

diff --git a/mushroom_rl/core/_impl/torch_dataset.py b/mushroom_rl/core/_impl/torch_dataset.py
@@ -149,6 +149,9 @@ def get_view(self, index):
             view._policy_states = self._policy_states[index, ...]
             view._policy_next_states = self._policy_next_states[index, ...]
 
+        if self._mask is not None:
+            view._mask = self._mask[index, ...]
+
         return view
 
     def __getitem__(self, index):