Hotfix viser traces (#107)

* Fixed viser and mujoco vis API issuse

* Added changes to changelog

* Added pycparser to deps

* Updated pixi lock

* Added John's more general changes to handle multi-colored splines

* Updated changelog and new package version

* Fixing the pixi lock file to be compatible with judo 0.0.6

Author: bhung-bdai

CHANGELOG.md

@@ -1,3 +1,13 @@
+# v0.0.6
+
+## Added
+* New copyrights
+
+## Fixed
+* Fixed bug that caused Viser and Mujoco to crash on visualization due to different API
+* Added pycparser deps
+* Fixed a new mujoco API compatibility
+
 # v0.0.5
 
 ## Added

judo/visualizers/model.py

@@ -211,7 +211,17 @@ def add_traces(
                 rgb=DEFAULT_BEST_SPLINE_COLOR,
             )
         )
-        self._traces[0].colors = np.tile(self._traces[0].colors[0, :, :], (all_traces_rollout_size, 1, 1))
+
+        # In viser 1.0.16+, colors can be shape (3,) if a single color was passed, or (N, 2, 3) if array was passed
+        # We need to handle both cases
+        if all_traces_rollout_size > 0:
+            if self._traces[0].colors.ndim == 1:
+                # Single color was stored, expand to (N, 2, 3) shape
+                self._traces[0].colors = np.tile(self._traces[0].colors, (all_traces_rollout_size, 2, 1))
+            else:
+                # Already (N, 2, 3), just tile the first segment
+                self._traces[0].colors = np.tile(self._traces[0].colors[:1, :, :], (all_traces_rollout_size, 1, 1))
+
         if (rest_trace_size := num_traces - all_traces_rollout_size) > 0:
             self._traces.append(
                 add_segments(
@@ -221,7 +231,12 @@ def add_traces(
                     rgb=DEFAULT_SPLINE_COLOR,
                 )
             )
-            self._traces[1].colors = np.tile(self._traces[1].colors[0, :, :], (rest_trace_size, 1, 1))
+            if self._traces[1].colors.ndim == 1:
+                # Single color was stored, expand to (N, 2, 3) shape
+                self._traces[1].colors = np.tile(self._traces[1].colors, (rest_trace_size, 2, 1))
+            else:
+                # Already (N, 2, 3), just tile the first segment
+                self._traces[1].colors = np.tile(self._traces[1].colors[:1, :, :], (rest_trace_size, 1, 1))
 
     def remove_traces(self) -> None:
         """Remove traces."""
@@ -236,7 +251,7 @@ def set_data(self, data: MjData) -> None:
             # Use atomic to update both position/orientation synchronously.
             with self._target.atomic():
                 # Line up order of bodies in spec with order of bodies in model.
-                data_idx = self._spec.bodies[i].id
+                data_idx = self._model.body(self._spec.bodies[i].name).id
                 self._bodies[i].position = tuple(data.xpos[data_idx])
                 self._bodies[i].wxyz = tuple(data.xquat[data_idx])
 

pixi.lock

@@ -5,7 +5,7 @@ build-backend = "hatchling.build"
 [project]
 name = "judo-rai"
 description = "Judo: a hackable sampling-based MPC toolbox"
-version = "0.0.5"
+version = "0.0.6"
 readme = "README.md"
 license = { file = "LICENSE" }
 authors = [
@@ -33,6 +33,7 @@ dependencies = [
     "mujoco",
     "numpy",
     "pillow",  # for displaying app logo
+    "pycparser",
     "requests",  # for sharing GUI link
     "rich",  # for displaying tables in benchmarking
     "scipy",