[BugFix] Fix bug with cached reset wrapper with counting the number of cached resets for user defined reset states (#1273)

Author: StoneT2000

mani_skill/utils/tree.py

@@ -22,3 +22,14 @@ def replace(x, i, y):
             replace(v, i, y[k])
     else:
         x[i] = y
+
+def shape(x, first_only=False):
+    """
+    Get the shape of leaf items in a tree. If first_only is True, return the shape of the first item only
+    """
+    if isinstance(x, dict):
+        if first_only:
+            return shape(next(iter(x.values())), first_only)
+        return {k: shape(v, first_only) for k, v in x.items()}
+    else:
+        return x.shape

mani_skill/utils/wrappers/cached_reset.py

@@ -58,7 +58,7 @@ def __init__(
         if reset_to_env_states is not None:
             self._cached_resets_env_states = reset_to_env_states["env_states"]
             self._cached_resets_obs_buffer = reset_to_env_states.get("obs", None)
-            self._num_cached_resets = len(self._cached_resets_env_states)
+            self._num_cached_resets = tree.shape(self._cached_resets_env_states)
         else:
             if self.cached_resets_config.num_resets is None:
                 self.cached_resets_config.num_resets = self.num_envs