[python/viewer] Refresh all nodes at once in meshcat.

python/jiminy_py/src/jiminy_py/viewer/meshcat/index.html

@@ -89,7 +89,9 @@
 
                         (async function() {
                             for await (const message of viewer.connection.messages) {
-                                viewer.handle_command_bytearray(new Uint8Array(message.buffers[0]));
+                                message.buffers.forEach(function(content, index) {
+                                    viewer.handle_command_bytearray(new Uint8Array(content));
+                                });
                             }
                             console.info("connection to Google Colab kernel closed.");
                         })();
@@ -100,7 +102,10 @@
                     viewer.connection = window.parent.Jupyter.notebook.kernel.comm_manager.new_comm("meshcat", "meshcat:open");
                     console.log("connection to Jupyter kernel:", viewer.connection);
                     viewer.connection.on_msg(function(message) {
-                        viewer.handle_command_bytearray(new Uint8Array(message.buffers[0].buffer));
+                        message.buffers.forEach(function(content, index) {
+                            console.log(content.buffer);
+                            viewer.handle_command_bytearray(new Uint8Array(content.buffer));
+                        });
                     });
                     viewer.connection.on_close(function(message) {
                         console.info("connection to Jupyter kernel closed:", message);
@@ -193,10 +198,12 @@
                             const reader = new FileReader();
                             reader.addEventListener('loadend', () => {
                                 var p = Promise.resolve(deserialize_array_buffer(reader.result));
-                                p.then(function(msg) {
-                                    if (msg.content.comm_id === comm_id)
+                                p.then(function(message) {
+                                    if (message.content.comm_id === comm_id)
                                     {
-                                        viewer.handle_command_bytearray(new Uint8Array(msg.buffers[0].buffer));
+                                        message.buffers.forEach(function(content, index) {
+                                            viewer.handle_command_bytearray(new Uint8Array(content.buffer));
+                                        });
                                     }
                                 });
                             });

python/jiminy_py/src/jiminy_py/viewer/meshcat/server.py

@@ -12,7 +12,7 @@
 import zmq
 from zmq.eventloop.zmqstream import ZMQStream
 
-from meshcat.servers.tree import walk
+from meshcat.servers.tree import walk, find_node
 from meshcat.servers.zmqserver import (
     DEFAULT_ZMQ_METHOD, VIEWER_ROOT, StaticFileHandlerNoCache,
     ZMQWebSocketBridge, WebSocketHandler, find_available_port)
@@ -180,12 +180,22 @@ def handle_zmq(self, frames: Sequence[bytes]) -> None:
                 websocket.write_message(msg, binary=True)
             for comm_id in self.comm_pool:
                 self.forward_to_comm(comm_id, msg)
+        elif cmd == "list":
+            # Only set_transform command is supported for now
+            for i in range(1, len(frames), 3):
+                cmd, path, data = frames[i:(i+3)]
+                path = list(filter(None, path.decode("utf-8").split("/")))
+                find_node(self.tree, path).transform = data
+                super().forward_to_websockets(frames[i:(i+3)])
+            for comm_id in self.comm_pool:
+                self.comm_zmq.send_multipart([comm_id, *frames[3::3]])
+            self.zmq_socket.send(b"ok")
         else:
             super().handle_zmq(frames)
 
     def handle_comm(self, frames: Sequence[bytes]) -> None:
         cmd = frames[0].decode("utf-8")
-        comm_id = f"{cmd.split(':', 2)[1]}".encode()
+        comm_id = cmd.split(':', 2)[1].encode()
         if cmd.startswith("open:"):
             self.send_scene(comm_id=comm_id)
             self.comm_pool.add(comm_id)
@@ -202,7 +212,7 @@ def handle_comm(self, frames: Sequence[bytes]) -> None:
             self.comm_msg.pop(comm_id, None)
         elif cmd.startswith("data:"):
             # Extract the message
-            message = f"{cmd.split(':', 2)[2]}"
+            message = cmd.split(':', 2)[2]
             # Catch watchdog messages
             if message == "watchdog":
                 self.watch_pool.add(comm_id)

python/jiminy_py/src/jiminy_py/viewer/meshcat/wrapper.py

@@ -58,13 +58,9 @@ def __init__(self):
             self.qsize_old = 0
             self.is_running = False
 
-        def __call__(self, unsafe: bool = False) -> None:
+        def __call__(self) -> None:
             """Check once if there is pending comm related event in the shell
             stream message priority queue.
-
-            :param unsafe: Whether or not to assume check if the number of
-                           pending message has changed is enough. It makes the
-                           evaluation much faster but flawed.
             """
             # Guard to avoid running this method several times in parallel
             if self.is_running:
@@ -91,15 +87,8 @@ def __call__(self, unsafe: bool = False) -> None:
                         shell_stream.socket, zmq.POLLIN)
                     events = shell_stream.poller.poll(0)
 
-            qsize = self.__kernel.msg_queue.qsize()
-            if unsafe and qsize == self.qsize_old:
-                # The number of queued messages in the queue has not changed
-                # since it last time it has been checked. Assuming those
-                # messages are the same has before and returning earlier.
-                qsize = 0
-
             # One must go through all the messages to keep them in order
-            for _ in range(qsize):
+            for _ in range(self.__kernel.msg_queue.qsize()):
                 *priority, t, dispatch, args = \
                     self.__kernel.msg_queue.get_nowait()
                 if not priority or priority[0] <= SHELL_PRIORITY:
@@ -200,19 +189,6 @@ def __call__(self, unsafe: bool = False) -> None:
     # Start comm hijacking
     process_kernel_comm = CommProcessor()
 
-    # Monkey-patch meshcat ViewerWindow 'send' method to process queued comm
-    # messages. Otherwise, new opening comm will not be detected soon enough.
-    _send_orig = meshcat.visualizer.ViewerWindow.send
-    def _send(self, command: Any) -> None:  # noqa
-        _send_orig(self, command)
-        # Check on new comm related messages. Unsafe is enabled to avoid
-        # potentially significant overhead. At this point several safe should
-        # have been executed, so it is much less likely than comm messages
-        # will slip through the net. Besides, missing messages at this point
-        # is not blocking, because here we are not waiting for it to continue.
-        process_kernel_comm(unsafe=True)
-    meshcat.visualizer.ViewerWindow.send = _send # noqa
-
 
 class CommManager:
     def __new__(cls, *args: Any, **kwargs: Any) -> "CommManager":
@@ -289,15 +265,11 @@ def close(self) -> None:
             self.__ioloop = None
 
     def __forward_to_ipykernel(self, frames: Sequence[bytes]) -> None:
-        try:
-            # There must be always two parts each messages, but who knows...
-            comm_id, cmd = frames
-        except ValueError:
-            return
+        comm_id, *cmd = frames
         comm_id = comm_id.decode()
         try:
             comm = self.__kernel.comm_manager.comms[comm_id]
-            comm.send(buffers=[cmd])
+            comm.send(buffers=cmd)
         except KeyError:
             # The comm has probably been closed without the server knowing.
             # Sending the notification to the server to consider it as such.

python/jiminy_py/src/jiminy_py/viewer/viewer.py

@@ -43,8 +43,7 @@
     Tuple3FType, Tuple4FType, Panda3dApp, Panda3dViewer, Panda3dVisualizer)
 
 
-REPLAY_FRAMERATE = 20 if interactive_mode() >= 3 else 30
-
+REPLAY_FRAMERATE = 30
 
 CAMERA_INV_TRANSFORM_PANDA3D = rpyToMatrix(-np.pi/2, 0.0, 0.0)
 CAMERA_INV_TRANSFORM_MESHCAT = rpyToMatrix(-np.pi/2, 0.0, 0.0)
@@ -152,6 +151,7 @@ def get_backend_type(backend_name: str) -> type:
             raise ImportError(
                 "'meshcat' backend is not available. Please install "
                 "'jiminy[meshcat]'.")
+    raise ValueError(f"Unknown backend '{backend_name}'.")
 
 
 def sleep(dt: float) -> None:
@@ -2023,15 +2023,30 @@ def refresh(self,
                     pose_dict[nodeName] = ((x, y, z), (qw, qx, qy, qz))
                 self._gui.move_nodes(group, pose_dict)
         else:
+            import umsgpack
             for geom_model, geom_data, model_type in zip(
                     model_list, data_list, model_type_list):
+                cmd_data = [b"list"]
                 for i, geom in enumerate(geom_model.geometryObjects):
                     oMg = geom_data.oMg[i]
                     S = np.diag((*geom.meshScale, 1.0))
-                    T = oMg.homogeneous.dot(S)
+                    H = oMg.homogeneous.dot(S)
                     nodeName = self._client.getViewerNodeName(
                         geom, model_type)
-                    self._gui[nodeName].set_transform(T)
+                    path = self._gui[nodeName].path.lower()
+                    cmd_data += [
+                        b"set_transform",
+                        path.encode("utf-8"),
+                        umsgpack.packb({
+                            "type": "set_transform",
+                            "path": path,
+                            "matrix": list(H.T.flat)
+                        })
+                    ]
+            # Moving all geometries at once using custom command to improve
+            # throughput due to piping sockets with multiple redirections.
+            self._gui.window.zmq_socket.send_multipart(cmd_data)
+            self._gui.window.zmq_socket.recv()
 
         # Update the camera placement if necessary
         if Viewer._camera_travelling is not None:

unit_py/test_flexible_arm.py

@@ -60,9 +60,9 @@ def generate_flexible_arm(mass_segments: float,
             izz=f"{inertia_segments}")
 
     motor = ET.SubElement(
-        robot, "joint", name=f"base_to_link0", type="revolute")
+        robot, "joint", name="base_to_link0", type="revolute")
     ET.SubElement(motor, "parent", link="base")
-    ET.SubElement(motor, "child", link=f"link0")
+    ET.SubElement(motor, "child", link="link0")
     ET.SubElement(motor, "origin", xyz="0 0 0", rpy=f"{np.pi/2} 0 0")
     ET.SubElement(motor, "axis", xyz="0 0 1")
     ET.SubElement(
@@ -89,7 +89,7 @@ class SimulateFlexibleArm(unittest.TestCase):
     def setUp(self):
         # Create temporary urdf file
         fd, urdf_path = tempfile.mkstemp(
-            prefix=f"flexible_arm_", suffix=".urdf")
+            prefix="flexible_arm_", suffix=".urdf")
         os.close(fd)
 
         # Procedural model generation
@@ -117,13 +117,13 @@ def setUp(self):
     def tearDown(self):
         Viewer.close()
 
-    def _read_write_replay_log(self, format):
+    def _read_write_replay_log(self, log_format):
         # Set initial condition and simulation duration
         q0, v0 = np.array([0.]), np.array([0.])
         t_end = 4.0
 
         # Generate temporary log file
-        ext = format if format != "binary" else "data"
+        ext = log_format if log_format != "binary" else "data"
         fd, log_path = tempfile.mkstemp(
             prefix=f"{self.simulator.robot.name}_", suffix=f".{ext}")
         os.close(fd)
@@ -135,7 +135,7 @@ def _read_write_replay_log(self, format):
 
         # Generate temporary video file
         fd, video_path = tempfile.mkstemp(
-            prefix=f"{self.simulator.robot.name}_", suffix=f".mp4")
+            prefix=f"{self.simulator.robot.name}_", suffix=".mp4")
         os.close(fd)
 
         # Record the result
@@ -173,8 +173,8 @@ def test_write_replay_standalone_log(self):
         self.simulator.robot.set_model_options(model_options)
 
         # Check both HDF5 and binary log formats
-        for format in ('hdf5', 'binary'):
-            self.assertTrue(self._read_write_replay_log(format))
+        for log_format in ('hdf5', 'binary'):
+            self.assertTrue(self._read_write_replay_log(log_format))
 
         # Make sure the scene is empty now
         self.assertEqual(len(Viewer._backend_robot_names), 0)
@@ -202,7 +202,7 @@ def test_rigid_vs_flex_at_frame(self):
 
         # Check different flexibility ordering
         q_flex, img_flex, pnc_model_flex, visual_model_flex = [], [], [], []
-        for ord in (range(N_FLEXIBILITY), range(N_FLEXIBILITY)[::-1]):
+        for order in (range(N_FLEXIBILITY), range(N_FLEXIBILITY)[::-1]):
             # Specify joint flexibility parameters
             model_options = self.simulator.robot.get_model_options()
             model_options['dynamics']['enableFlexibleModel'] = True
@@ -211,7 +211,7 @@ def test_rigid_vs_flex_at_frame(self):
                 'stiffness': np.zeros(3),
                 'damping': np.zeros(3),
                 'inertia': np.full(3, fill_value=1e6)
-            } for i in ord]
+            } for i in order]
             self.simulator.robot.set_model_options(model_options)
 
             # Serialize the model

unit_py/utilities.py

@@ -2,8 +2,6 @@
 @brief  Utility functions for unit tests.
 """
 import os
-import shutil
-import tempfile
 import numpy as np
 from scipy.integrate import ode
 from typing import Optional, Union, Dict, Sequence, Tuple, Callable