fix: shuffle demos before slicing to get different demos

bigym/bigym_env.py

@@ -5,19 +5,19 @@
 from pathlib import Path
 from typing import Any, Optional, Type
 
+import gymnasium as gym
 import mujoco
 import numpy as np
-import gymnasium as gym
 from gymnasium import spaces
 from mojo import Mojo
-from mojo.elements import Geom, Camera
+from mojo.elements import Camera, Geom
 
 from bigym.action_modes import ActionMode
+from bigym.bigym_renderer import BiGymRenderer
 from bigym.const import WORLD_MODEL
 from bigym.envs.props.preset import Preset
 from bigym.robots.configs.h1 import H1
 from bigym.robots.robot import Robot
-from bigym.bigym_renderer import BiGymRenderer
 from bigym.utils.callables_cache import CallablesCache
 from bigym.utils.env_health import EnvHealth
 from bigym.utils.observation_config import ObservationConfig

demonstrations/demo_store.py

@@ -1,24 +1,23 @@
 """Script for uploading the collected demos."""
 import logging
 import os
+import tempfile
 import warnings
 import zipfile
+from copy import deepcopy
+from pathlib import Path
 from typing import Optional
 
 import numpy as np
-import tempfile
-from pathlib import Path
-from copy import deepcopy
-
 import wget
 from tqdm import tqdm
 
 from bigym.bigym_env import CONTROL_FREQUENCY_MAX
 from bigym.const import CACHE_PATH, DEMO_RELEASES, DEMO_VERSION
 from demonstrations.const import SAFETENSORS_SUFFIX
-from demonstrations.utils import Metadata, ObservationMode
 from demonstrations.demo import Demo, LightweightDemo
 from demonstrations.demo_converter import DemoConverter
+from demonstrations.utils import Metadata, ObservationMode
 
 
 class DemoNotFoundError(Exception):
@@ -165,9 +164,10 @@ def _get_demos(self, demos_dir: Path, amount: int) -> list[Demo]:
         files = list(demos_dir.glob(f"*{SAFETENSORS_SUFFIX}"))
         if amount > len(files):
             raise TooManyDemosRequestedError(amount, len(files))
-        elif amount > 0:
-            files = files[:amount]
+
         np.random.shuffle(files)
+        if amount > 0:
+            files = files[:amount]
         return [Demo.from_safetensors(file) for file in files]
 
     def _get_demos_count(self, demos_dir: Path) -> int: