Fix bug in privileged

Author: yardenas

ss2r/benchmark_suites/rccar/rccar.py

@@ -45,7 +45,7 @@ def randomize(rng):
     cfg = OmegaConf.to_container(cfg)
     in_axes = jax.tree_map(lambda _: 0, sys)
     sys, params = randomize(rng)
-    return sys, in_axes, params[:, None]
+    return sys, in_axes, params
 
 
 def rotate_coordinates(state: jnp.array, encode_angle: bool = False) -> jnp.array:

ss2r/common/logging.py

@@ -118,7 +118,10 @@ def log_video(
         fps: int | float = 30,
         flush: bool = False,
     ):
-        self._writer.add_video(name, np.array(images, copy=False), step, fps=fps)
+        images = np.array(images, copy=False)
+        if images.ndim == 4:
+            images = images[None]
+        self._writer.add_video(name, images, step, fps=fps)
         if flush:
             self._writer.flush()