Simplify implementation of 2D plotting

Author: rohanpsingh

mujoco_viewer/mujoco_viewer.py

@@ -73,14 +73,8 @@ def __init__(
         self.fig = mujoco.MjvFigure()
         mujoco.mjv_defaultFigure(self.fig)
 
-        # Points for sampling of sensors... dictates smoothness of graph
-        self._num_pnts = 100
-        self._data_graph_line_names = []
-        self._line_datas = []
-
-        for n in range(0, len(self.model.sensor_adr) * 3):
-            for i in range(0, 300):
-                self.fig.linedata[n][2 * i] = float(-i)
+        # Counter for number of lines added to plot
+        self._line_counter = 0
 
         self.ctx = mujoco.MjrContext(
             self.model, mujoco.mjtFontScale.mjFONTSCALE_150.value)
@@ -98,9 +92,6 @@ def __init__(
         self.fig.flg_extend = 1
         self.fig.flg_symmetric = 0
 
-        # Makes the graph to be in autorange
-        self.axis_autorange()
-
         # load camera from configuration (if available)
         pathlib.Path(
             self.CONFIG_PATH.parent).mkdir(
@@ -146,89 +137,41 @@ def __init__(
         self._overlay = {}
         self._markers = []
 
-    def set_grid_divisions(self, x_div: int, y_div: int, x_axis_time: float = 0.0, override=False):
-        if override is False:
-            assert x_axis_time >= self.model.opt.timestep * 50, "Set [x_axis_time] >= [self.model.opt.timestep * 50], inorder to get a suitable sampling rate"
-        self.fig.gridsize[0] = x_div + 1
-        self.fig.gridsize[1] = y_div + 1
-        if x_axis_time != 0.0:
-            self._num_pnts = x_axis_time / self.model.opt.timestep
-            print("self._num_pnts: ", self._num_pnts)
-            if self._num_pnts > 300:
-                self._num_pnts = 300
-                new_x_axis_time = self.model.opt.timestep * self._num_pnts
-                print(
-                    f"Minimum x_axis_time is: {new_x_axis_time}"
-                    + " reduce the x_axis_time"
-                    f" OR Maximum time_step is: "
-                    + f"{self.model.opt.timestep*self._num_pnts}"
-                    + " increase the timestep"
-                )
-                # assert x_axis_time ==
-            assert 1 <= self._num_pnts <= 300, (
-                "num_pnts should be [10,300], it is currently:",
-                f"{self._num_pnts}",
-            )
-            # self._num_pnts = num_pnts
-            self._time_per_div = (self.model.opt.timestep * self._num_pnts) / (
-                x_div
-            )
-            self.set_x_label(
-                xname=f"time/div: {self._time_per_div}s"
-                + f" total: {self.model.opt.timestep * self._num_pnts}"
-            )
-
-    def axis_autorange(self):
-        """
-        Call this function to auto-range the graph
-        """
-        self.fig.range[0][0] = 1.0
-        self.fig.range[0][1] = -1.0
-        self.fig.range[1][0] = 1.0
-        self.fig.range[1][1] = -1.0
-
-    def set_graph_name(self, name: str):
-        assert type(name) == str, "name is not a string"
-        self.fig.title = name
-
-    def show_graph_legend(self, show_legend: bool = True):
-        if show_legend is True:
-            for i in range(0, len(self._data_graph_line_names)):
-                self.fig.linename[i] = self._data_graph_line_names[i]
-            self.fig.flg_legend = True
-
-    def set_x_label(self, xname: str):
-        assert type(xname) == str, "xname is not a string"
-        self.fig.xlabel = xname
-
-    def add_graph_line(self, line_name, line_data=0.0):
+    def add_graph_line(self, line_name):
         assert (
             type(line_name) == str
         ), f"Line_name is not a string: {type(line_name)}"
-        if line_name in self._data_graph_line_names:
-            print("line name already exists")
-        else:
-            self._data_graph_line_names.append(line_name)
-            self._line_datas.append(line_data)
+        if line_name in []:
+            raise Exception(
+                "line name already exists"
+            )
+
+        self.fig.linename[self._line_counter] = line_name
+        for i in range(mujoco.mjMAXLINEPNT):
+            self.fig.linedata[self._line_counter][2*i] = float(i)
+        self.fig.linepnt[self._line_counter] = 0
+        self._line_counter+=1
 
     def update_graph_line(self, line_name, line_data):
-        if line_name in self._data_graph_line_names:
-            idx = self._data_graph_line_names.index(line_name)
-            self._line_datas[idx] = line_data
-        else:
-            raise NameError(
+        try:
+            _line_name = line_name.encode('utf8')
+            linenames = self.fig.linename.tolist()
+            line_idx = linenames.index(_line_name)
+        except ValueError:
+            raise Exception(
                 "line name is not valid, add it to list before calling update"
             )
 
-    def sensorupdate(self):
-        pnt = int(mujoco.mju_min(self._num_pnts, self.fig.linepnt[0] + 1))
-        for n in range(0, len(self._line_datas)):
-            for i in range(pnt - 1, 0, -1):
-                self.fig.linedata[n][2 * i + 1] = self.fig.linedata[n][
-                    2 * i - 1
-                ]
-            self.fig.linepnt[n] = pnt
-            self.fig.linedata[n][1] = self._line_datas[n]
+        _x = self.fig.linepnt[line_idx]
+
+        if _x < mujoco.mjMAXLINEPNT:
+            self.fig.linedata[line_idx][2*_x] = _x
+            self.fig.linedata[line_idx][2*_x + 1] = float(line_data)
+            self.fig.linepnt[line_idx] += 1
+        else:
+            for i in range(mujoco.mjMAXLINEPNT-1):
+                self.fig.linedata[line_idx][2*i+1] = self.fig.linedata[line_idx][2*(i+1)+1]
+            self.fig.linedata[line_idx][-1] = float(line_data)
 
     def update_graph_size(self, size_div_x=None, size_div_y=None):
         if size_div_x is None and size_div_y is None:
@@ -429,7 +372,6 @@ def read_pixels(self, camid=None, depth=False):
         mujoco.mjr_render(self.viewport, self.scn, self.ctx)
         shape = glfw.get_framebuffer_size(self.window)
 
-        
         if depth:
             rgb_img = np.zeros((shape[1], shape[0], 3), dtype=np.uint8)
             depth_img = np.zeros((shape[1], shape[0], 1), dtype=np.float32)
@@ -491,24 +433,11 @@ def update():
                         self.ctx)
 
                 # Handle graph and pausing interactions
-                if (
-                    not self._paused
-                    and not self._hide_graph
-                ):
-                    self.sensorupdate()
-                    self.update_graph_size()
-                    mujoco.mjr_figure(
-                        self.graph_viewport, self.fig, self.ctx
-                    )
-                elif self._hide_graph and self._paused:
-                    self.update_graph_size()
-                elif not self._hide_graph and self._paused:
+                if not self._hide_graph:
                     mujoco.mjr_figure(
                         self.graph_viewport, self.fig, self.ctx
                     )
-                elif self._hide_graph and not self._paused:
-                    self.sensorupdate()
-                    self.update_graph_size()
+                self.update_graph_size()
 
                 glfw.swap_buffers(self.window)
             glfw.poll_events()