feat: extend reward plot for multiple signals

.github/workflows/main.yml

@@ -15,7 +15,7 @@ jobs:
     runs-on: ubuntu-latest
     strategy:
       matrix:
-        python-version: ["3.8", "3.9", "3.10"]
+        python-version: ["3.8", "3.9", "3.10", "3.11", "3.12"]
     steps:
     - uses: actions/checkout@v3
     - name: Set up Python ${{ matrix.python-version }}

.github/workflows/pull_request.yml

@@ -9,7 +9,7 @@ jobs:
     runs-on: ubuntu-latest
     strategy:
       matrix:
-        python-version: ["3.8", "3.9", "3.10"]
+        python-version: ["3.8", "3.9", "3.10", "3.11", "3.12"]
     steps:
     - uses: actions/checkout@v3
     - name: Set up Python ${{ matrix.python-version }}

pyproject.toml

@@ -14,7 +14,7 @@ include = [
 name = "dm-robotics-panda"
 description = "Panda model for dm_robotics."
 version = "0.4.7"
-requires-python = ">=3.8,<3.11"
+requires-python = ">=3.8,<3.13"
 authors = [
     { name = "Jean Elsner", email = "jean.elsner@tum.de" },
 ]
@@ -40,6 +40,8 @@ classifiers = [
         "Programming Language :: Python :: 3.8",
         "Programming Language :: Python :: 3.9",
         "Programming Language :: Python :: 3.10",
+        "Programming Language :: Python :: 3.11",
+        "Programming Language :: Python :: 3.12",
 ]
 
 [project.optional-dependencies]

src/dm_robotics/panda/parameters.py

@@ -64,5 +64,6 @@ class RobotParams:
   robot_ip: Optional[str] = None
   joint_stiffness: Sequence[float] = (600, 600, 600, 600, 250, 150, 50)
   joint_damping: Sequence[float] = (50, 50, 50, 20, 20, 20, 10)
-  collision_behavior: CollisionBehavior = CollisionBehavior()
+  collision_behavior: CollisionBehavior = dataclasses.field(
+      default_factory=lambda: CollisionBehavior())
   enforce_realtime: bool = False

src/dm_robotics/panda/utils.py

@@ -256,19 +256,35 @@
                maxlen: int = 500) -> None:
     super().__init__(runtime, maxlen)
     self.fig.title = 'Reward'
-    self.maxlines = 1
-    self.y.append(deque(maxlen=self.maxlen))
+    self.maxlines = None
+
+  def _init_buffer(self):
+    if isinstance(self._rt._time_step.reward, np.ndarray):
+      self.maxlines = self._rt._time_step.reward.shape[0]
+    else:
+      self.maxlines = 1
+    for _1 in range(self.maxlines):
+      self.y.append(deque(maxlen=self.maxlen))
     self.reset_data()
 
   def render(self, context, viewport):
-    if self._rt._time_step is None:
+    if self._rt._time_step is None or self._rt._time_step.reward is None:
       return
-    r = self._rt._time_step.reward
-    self.fig.linepnt[0] = self.maxlen
-    self.y[0].append(r)
-    self.fig.linedata[0][:self.maxlen * 2] = np.array([self.x,
-                                                       self.y[0]]).T.reshape(
-                                                           (-1,))
+    if self.maxlines is None:
+      self._init_buffer()
+    if self.maxlines > 1:
+      for i, r in enumerate(self._rt._time_step.reward):
+        self.fig.linepnt[i] = self.maxlen
+        self.y[i].append(r)
+        self.fig.linedata[i][:self.maxlen * 2] = np.array([self.x, self.y[i]
+                                                          ]).T.reshape((-1,))
+    else:
+      r = self._rt._time_step.reward
+      self.fig.linepnt[0] = self.maxlen
+      self.y[0].append(r)
+      self.fig.linedata[0][:self.maxlen * 2] = np.array([self.x,
+                                                         self.y[0]]).T.reshape(
+                                                             (-1,))
     pos = mujoco.MjrRect(2 * 300 + 5, viewport.height - 200 - 5, 300, 200)
     mujoco.mjr_figure(pos, self.fig, context.ptr)