From 72481508860182640f9f7bec7b012e35c0e557a9 Mon Sep 17 00:00:00 2001
From: Alexis DUBURCQ <alexis.duburcq@gmail.com>
Date: Mon, 7 Dec 2020 21:31:14 +0100
Subject: [PATCH] [python|gym] Performance improvement. (#242)

* [core/python] Do not use Eigenpy to convert systemState.q/v/a to improve efficiency and preserve constness.
* [python/viewer] Enable to add legend in Meshcat.
* [python/viewer] Enable to add logo in meshcat.
* [python/viewer] Add 'legend' and 'logo_fullpath' optional argument to 'play_trajectories'.
* [python|gym] Add Python native 'is_simulation_running' attribute to Simulator for fast access.
* [gym] Define shared memories/proxies at init when possible.
* [gym] Replace 'engine' proxy by 'simulator' for greater flexibility.
* [gym/envs] Improve performance.

Co-authored-by: Alexis Duburcq <alexis.duburcq@wandercraft.eu>
---
 CMakeLists.txt                                |  2 +-
 .../gym_jiminy/common/bases/generic_bases.py  |  5 +-
 .../gym_jiminy/common/bases/pipeline_bases.py | 25 +++---
 .../gym_jiminy/common/envs/env_generic.py     | 35 ++++----
 .../gym_jiminy/common/envs/env_locomotion.py  | 10 +--
 .../common/gym_jiminy/common/pipeline.py      |  2 +-
 .../gym_jiminy/common/wrappers/frame_stack.py |  5 +-
 .../envs/gym_jiminy/envs/acrobot.py           |  9 ++-
 .../envs/gym_jiminy/envs/cartpole.py          | 10 ++-
 python/jiminy_py/src/jiminy_py/log.py         |  2 +-
 .../src/jiminy_py/meshcat/index.html          | 33 ++++++++
 .../jiminy_py/src/jiminy_py/meshcat/legend.js | 79 +++++++++++++++++++
 .../jiminy_py/src/jiminy_py/meshcat/server.py |  9 +--
 .../src/jiminy_py/meshcat/wrapper.py          | 62 +++++++++++++++
 python/jiminy_py/src/jiminy_py/simulator.py   | 45 +++++++++--
 python/jiminy_py/src/jiminy_py/viewer.py      | 55 ++++++++++++-
 .../include/jiminy/python/Jiminy.h            | 27 +++++--
 17 files changed, 341 insertions(+), 74 deletions(-)
 create mode 100644 python/jiminy_py/src/jiminy_py/meshcat/legend.js

diff --git a/CMakeLists.txt b/CMakeLists.txt
index 29ede5df8..e2bf0c4c8 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -2,7 +2,7 @@
 cmake_minimum_required(VERSION 3.10)
 
 # Set the build version
-set(BUILD_VERSION 1.4.21)
+set(BUILD_VERSION 1.4.22)
 
 # Add definition of Jiminy version for C++ headers
 add_definitions("-DJIMINY_VERSION=\"${BUILD_VERSION}\"")
diff --git a/python/gym_jiminy/common/gym_jiminy/common/bases/generic_bases.py b/python/gym_jiminy/common/gym_jiminy/common/bases/generic_bases.py
index 390427c11..8cd1cca68 100644
--- a/python/gym_jiminy/common/gym_jiminy/common/bases/generic_bases.py
+++ b/python/gym_jiminy/common/gym_jiminy/common/bases/generic_bases.py
@@ -6,6 +6,7 @@
 import gym
 
 import jiminy_py.core as jiminy
+from jiminy_py.simulator import Simulator
 
 from ..utils import copy, SpaceDictNested
 
@@ -173,14 +174,14 @@ class ObserverControllerInterface(ObserverInterface, ControllerInterface):
     """Observer plus controller interface for both generic pipeline blocks,
     including environments.
     """
-    engine: Optional[jiminy.EngineMultiRobot]
+    simulator: Optional[Simulator]
     stepper_state: Optional[jiminy.StepperState]
     system_state: Optional[jiminy.SystemState]
     sensors_data: Optional[Dict[str, np.ndarray]]
 
     def __init__(self, *args: Any, **kwargs: Any) -> None:
         # Define some attributes
-        self.engine = None
+        self.simulator = None
         self.stepper_state = None
         self.system_state = None
         self.sensors_data = None
diff --git a/python/gym_jiminy/common/gym_jiminy/common/bases/pipeline_bases.py b/python/gym_jiminy/common/gym_jiminy/common/bases/pipeline_bases.py
index 0e8043693..32e860884 100644
--- a/python/gym_jiminy/common/gym_jiminy/common/bases/pipeline_bases.py
+++ b/python/gym_jiminy/common/gym_jiminy/common/bases/pipeline_bases.py
@@ -18,6 +18,7 @@
 import gym
 
 import jiminy_py.core as jiminy
+from jiminy_py.simulator import Simulator
 from jiminy_py.controller import ObserverHandleType, ControllerHandleType
 
 from ..utils import (
@@ -54,6 +55,12 @@ def __init__(self,
         # Initialize base wrapper and interfaces through multiple inheritance
         super().__init__(env)  # Do not forward extra arguments, if any
 
+        # Refresh some proxies for fast lookup
+        self.simulator: Simulator = self.env.simulator
+        self.stepper_state: jiminy.StepperState = self.env.stepper_state
+        self.system_state: jiminy.SystemState = self.env.system_state
+        self.sensors_data: jiminy.sensorsData = self.env.sensors_data
+
         # Define some internal buffers
         self._dt_eps: Optional[float] = None
         self._command = zeros(self.env.unwrapped.action_space)
@@ -187,9 +194,6 @@ def _setup(self) -> None:
         fill(self._command, 0.0)
 
         # Refresh some proxies for fast lookup
-        self.engine = self.env.engine
-        self.stepper_state = self.env.stepper_state
-        self.system_state = self.env.system_state
         self.sensors_data = self.env.sensors_data
 
     def refresh_observation(self) -> None:  # type: ignore[override]
@@ -344,9 +348,6 @@ def refresh_observation(self) -> None:  # type: ignore[override]
         """
         # pylint: disable=arguments-differ
 
-        # Assertion(s) for type checker
-        assert self.engine is not None and self.stepper_state is not None
-
         # Get environment observation
         super().refresh_observation()
 
@@ -355,7 +356,7 @@ def refresh_observation(self) -> None:  # type: ignore[override]
         if _is_breakpoint(t, self.observe_dt, self._dt_eps):
             obs = self.env.get_observation()
             self.observer.refresh_observation(obs)
-            if not self.engine.is_simulation_running:
+            if not self.simulator.is_simulation_running:
                 features = self.observer.get_observation()
                 if self.augment_observation:
                     self._observation = obs
@@ -532,9 +533,6 @@ def compute_command(self,
         :param measure: Observation of the environment.
         :param action: High-level target to achieve.
         """
-        # Assertion(s) for type checker
-        assert self.engine is not None and self.stepper_state is not None
-
         # Backup the action
         set_value(self._action, action)
 
@@ -552,7 +550,7 @@ def compute_command(self,
         # update the command of the right period. Ultimately, this is done
         # automatically by the engine, which is calling `_controller_handle` at
         # the right period.
-        if self.engine.is_simulation_running:
+        if self.simulator.is_simulation_running:
             # Do not update command during the first iteration because the
             # action is undefined at this point
             np.core.umath.copyto(self._command, self.env.compute_command(
@@ -576,14 +574,11 @@ def refresh_observation(self) -> None:  # type: ignore[override]
         :returns: Original environment observation, eventually including
                   controllers targets if requested.
         """
-        # Assertion(s) for type checker
-        assert self.engine is not None
-
         # Get environment observation
         super().refresh_observation()
 
         # Add target to observation if requested
-        if not self.engine.is_simulation_running:
+        if not self.simulator.is_simulation_running:
             self._observation = self.env.get_observation()
             if self.augment_observation:
                 self._observation.setdefault('targets', {})[
diff --git a/python/gym_jiminy/common/gym_jiminy/common/envs/env_generic.py b/python/gym_jiminy/common/gym_jiminy/common/envs/env_generic.py
index ecddd1552..4de019df7 100644
--- a/python/gym_jiminy/common/gym_jiminy/common/envs/env_generic.py
+++ b/python/gym_jiminy/common/gym_jiminy/common/envs/env_generic.py
@@ -90,14 +90,20 @@ def __init__(self,
         """
         # pylint: disable=unused-argument
 
+        # Initialize the interfaces through multiple inheritance
+        super().__init__()  # Do not forward extra arguments, if any
+
         # Backup some user arguments
-        self.simulator = simulator
+        self.simulator: Simulator = simulator
         self.step_dt = step_dt
         self.enforce_bounded = enforce_bounded
         self.debug = debug
 
-        # Initialize the interfaces through multiple inheritance
-        super().__init__()  # Do not forward extra arguments, if any
+        # Define some proxies for fast access
+        self.engine: jiminy.EngineMultiRobot = self.simulator.engine
+        self.stepper_state: jiminy.StepperState = self.engine.stepper_state
+        self.system_state: jiminy.SystemState = self.engine.system_state
+        self.sensors_data: jiminy.sensorsData = dict(self.robot.sensors_data)
 
         # Internal buffers for physics computations
         self.rg = np.random.RandomState()
@@ -418,11 +424,8 @@ def reset(self,
         # Reset the simulator
         self.simulator.reset()
 
-        # Initialize shared memories.
+        # Re-initialize some shared memories.
         # It must be done because the robot may have changed.
-        self.engine = self.simulator.engine
-        self.stepper_state = self.engine.stepper_state
-        self.system_state = self.engine.system_state
         self.sensors_data = dict(self.robot.sensors_data)
 
         # Make sure the environment is properly setup
@@ -503,14 +506,14 @@ def reset(self,
             is_obs_valid = False
         if not is_obs_valid:
             raise RuntimeError(
-                "The observation returned by `refresh_observation` is "
+                "The observation computed by `refresh_observation` is "
                 "inconsistent with the observation space defined by "
                 "`_refresh_observation_space`.")
 
         if self.is_done():
             raise RuntimeError(
-                "The simulation is already done at `reset`. "
-                "Check the implementation of `is_done` if overloaded.")
+                "The simulation is already done at `reset`. Check the "
+                "implementation of `is_done` if overloaded.")
 
         return self.get_observation()
 
@@ -558,7 +561,7 @@ def step(self,
                   not), and a dictionary of extra information
         """
         # Make sure a simulation is already running
-        if self.engine is None or not self.engine.is_simulation_running:
+        if not self.simulator.is_simulation_running:
             raise RuntimeError(
                 "No simulation running. Please call `reset` before `step`.")
 
@@ -625,10 +628,9 @@ def step(self,
     def get_log(self) -> Tuple[Dict[str, np.ndarray], Dict[str, str]]:
         """Get log of recorded variable since the beginning of the episode.
         """
-        if self.engine is None or not self.engine.is_simulation_running:
+        if not self.simulator.is_simulation_running:
             raise RuntimeError(
-                "No simulation running. Please call `reset` at least one "
-                "before getting log.")
+                "No simulation running. Please start one before getting log.")
         return self.simulator.get_log()
 
     def render(self,
@@ -839,11 +841,10 @@ def refresh_observation(self) -> None:  # type: ignore[override]
         # pylint: disable=arguments-differ
 
         # Assertion(s) for type checker
-        assert (self.engine is not None and self.stepper_state is not None and
-                isinstance(self._observation, dict))
+        assert isinstance(self._observation, dict)
 
         self._observation['t'][0] = self.stepper_state.t
-        if not self.engine.is_simulation_running:
+        if not self.simulator.is_simulation_running:
             (self._observation['state']['Q'],
              self._observation['state']['V']) = self.simulator.state
             self._observation['sensors'] = self.sensors_data
diff --git a/python/gym_jiminy/common/gym_jiminy/common/envs/env_locomotion.py b/python/gym_jiminy/common/gym_jiminy/common/envs/env_locomotion.py
index d3fc78b8c..6e7170c9c 100644
--- a/python/gym_jiminy/common/gym_jiminy/common/envs/env_locomotion.py
+++ b/python/gym_jiminy/common/gym_jiminy/common/envs/env_locomotion.py
@@ -384,9 +384,10 @@ def is_done(self) -> bool:  # type: ignore[override]
         """
         # pylint: disable=arguments-differ
 
-        # Assertion(s) for type checker
-        assert self.system_state is not None
-
+        if not self.simulator.is_simulation_running:
+            raise RuntimeError(
+                "No simulation running. Please start one before calling this "
+                "method.")
         if self.system_state.q[2] < self._height_neutral * 0.75:
             return True
         if self.simulator.stepper_state.t >= self.simu_duration_max:
@@ -408,9 +409,6 @@ def compute_reward(self,  # type: ignore[override]
         """
         # pylint: disable=arguments-differ
 
-        # Assertion(s) for type checker
-        assert self.system_state is not None
-
         reward_dict = info.setdefault('reward', {})
 
         # Define some proxies
diff --git a/python/gym_jiminy/common/gym_jiminy/common/pipeline.py b/python/gym_jiminy/common/gym_jiminy/common/pipeline.py
index efa0224ec..c07d27824 100644
--- a/python/gym_jiminy/common/gym_jiminy/common/pipeline.py
+++ b/python/gym_jiminy/common/gym_jiminy/common/pipeline.py
@@ -234,7 +234,7 @@ def __init__(self: wrapped_env_class,  # type: ignore[valid-type]
 def load_pipeline(fullpath: str) -> Type[BasePipelineWrapper]:
     """ TODO: Write documentation.
     """
-    file_ext = ''.join(pathlib.Path(fullpath).suffixes)
+    file_ext = pathlib.Path(fullpath).suffix
     with open(fullpath, 'r') as f:
         if file_ext == '.json':
             return build_pipeline(**json.load(f))
diff --git a/python/gym_jiminy/common/gym_jiminy/common/wrappers/frame_stack.py b/python/gym_jiminy/common/gym_jiminy/common/wrappers/frame_stack.py
index 663eb075a..0d61baa91 100644
--- a/python/gym_jiminy/common/gym_jiminy/common/wrappers/frame_stack.py
+++ b/python/gym_jiminy/common/gym_jiminy/common/wrappers/frame_stack.py
@@ -182,15 +182,12 @@ def _setup(self) -> None:
                 "`StackedJiminyEnv` does not support time-continuous update.")
 
     def refresh_observation(self) -> None:  # type: ignore[override]
-        # Assertion(s) for type checker
-        assert self.engine is not None and self.stepper_state is not None
-
         # Get environment observation
         self.env.refresh_observation()
 
         # Update observed features if necessary
         t = self.stepper_state.t
-        if self.engine.is_simulation_running and \
+        if self.simulator.is_simulation_running and \
                 _is_breakpoint(t, self.observe_dt, self._dt_eps):
             self.__n_last_stack += 1
         if self.__n_last_stack == self.skip_frames_ratio:
diff --git a/python/gym_jiminy/envs/gym_jiminy/envs/acrobot.py b/python/gym_jiminy/envs/gym_jiminy/envs/acrobot.py
index ba6349c11..1d81fd299 100644
--- a/python/gym_jiminy/envs/gym_jiminy/envs/acrobot.py
+++ b/python/gym_jiminy/envs/gym_jiminy/envs/acrobot.py
@@ -107,6 +107,10 @@ def __init__(self, continuous: bool = False):
         # Configure the learning environment
         super().__init__(simulator, STEP_DT, debug=False)
 
+        # Create some proxies for fast access
+        self.__state_view = (self._observation[:self.robot.nq],
+                             self._observation[self.robot.nv:])
+
     def _refresh_observation_space(self) -> None:
         """Configure the observation of the environment.
 
@@ -131,7 +135,10 @@ def refresh_observation(self, *args: Any, **kwargs: Any) -> None:
             For goal env, in addition of the current robot state, both the
             desired and achieved goals are observable.
         """
-        np.concatenate(self.simulator.state, out=self._observation)
+        if not self.simulator.is_simulation_running:
+            self.__state = (self.system_state.q, self.system_state.v)
+        np.core.umath.copyto(self.__state_view[0], self.__state[0])
+        np.core.umath.copyto(self.__state_view[1], self.__state[1])
 
     def _refresh_action_space(self) -> None:
         """Configure the action space of the environment.
diff --git a/python/gym_jiminy/envs/gym_jiminy/envs/cartpole.py b/python/gym_jiminy/envs/gym_jiminy/envs/cartpole.py
index 32ec80205..e37863ded 100644
--- a/python/gym_jiminy/envs/gym_jiminy/envs/cartpole.py
+++ b/python/gym_jiminy/envs/gym_jiminy/envs/cartpole.py
@@ -123,6 +123,10 @@ def __init__(self, continuous: bool = False):
         # Configure the learning environment
         super().__init__(simulator, STEP_DT, debug=False)
 
+        # Create some proxies for fast access
+        self.__state_view = (self._observation[:self.robot.nq],
+                             self._observation[self.robot.nv:])
+
     def _setup(self) -> None:
         """ TODO: Write documentation.
         """
@@ -176,8 +180,10 @@ def _sample_state(self) -> Tuple[np.ndarray, np.ndarray]:
 
     def refresh_observation(self) -> None:
         # @copydoc BaseJiminyEnv::refresh_observation
-        np.core.umath.copyto(
-            self._observation, np.concatenate(self.simulator.state))
+        if not self.simulator.is_simulation_running:
+            self.__state = (self.system_state.q, self.system_state.v)
+        np.core.umath.copyto(self.__state_view[0], self.__state[0])
+        np.core.umath.copyto(self.__state_view[1], self.__state[1])
 
     def is_done(self) -> bool:
         """ TODO: Write documentation.
diff --git a/python/jiminy_py/src/jiminy_py/log.py b/python/jiminy_py/src/jiminy_py/log.py
index 2093719d8..571bd36ed 100644
--- a/python/jiminy_py/src/jiminy_py/log.py
+++ b/python/jiminy_py/src/jiminy_py/log.py
@@ -48,7 +48,7 @@ def read_log(fullpath: str,
     """
     # Handling of file file_format
     if file_format is None:
-        file_ext = ''.join(pathlib.Path(fullpath).suffixes)
+        file_ext = pathlib.Path(fullpath).suffix
         if file_ext == '.data':
             file_format = 'binary'
         elif file_ext == '.csv' or file_ext == '.txt':
diff --git a/python/jiminy_py/src/jiminy_py/meshcat/index.html b/python/jiminy_py/src/jiminy_py/meshcat/index.html
index 661ce0a46..dcc85bdb7 100644
--- a/python/jiminy_py/src/jiminy_py/meshcat/index.html
+++ b/python/jiminy_py/src/jiminy_py/meshcat/index.html
@@ -10,6 +10,7 @@
 
         <script type="text/javascript" src="main.min.js"></script>
         <script type="text/javascript" src="webm-writer-0.3.0.js"></script>
+        <script type="text/javascript" src="legend.js"></script>
         <script>
             // Instantiate a new Meshcat viewer
             var viewer = new MeshCat.Viewer(document.getElementById("meshcat-pane"), false);
@@ -22,6 +23,38 @@
             viewer.handle_command = function(cmd) {
                 if (cmd.type == "ready") {
                     viewer.connection.send("meshcat:ok");
+                } else if (cmd.type == "legend") {
+                    var legend = document.getElementById("legend");
+                    if (legend == null) {
+                        createLegend("legend");
+                        legend = document.getElementById("legend");
+                    }
+                    if (cmd.text) {
+                        setLegendItem(legend, cmd.id, cmd.text, cmd.color);
+                    } else {
+                        removeLegendItem(legend, cmd.id);
+                    }
+                } else if (cmd.type == "logo") {
+                    var logo = document.getElementById("logo");
+                    if (cmd.data) {
+                        if (logo == null) {
+                            logo = document.createElement("img");
+                            logo.id = "logo";
+                            logo.draggable = false;
+                            logo.style.position = "fixed";
+                            logo.style.bottom = "20px";
+                            logo.style.left = "20px";
+                            logo.style.pointerEvents = "none";
+                            document.body.prepend(logo);
+                        }
+                        logo.setAttribute('src', 'data:image/png;base64,' + cmd.data);
+                        logo.style.width = cmd.width.toString() + "px";
+                        logo.style.height = cmd.height.toString() + "px";
+                    } else {
+                        if (logo !== null) {
+                            document.body.removeChild(logo);
+                        }
+                    }
                 } else {
                     handle_command.call(this, cmd);
                 }
diff --git a/python/jiminy_py/src/jiminy_py/meshcat/legend.js b/python/jiminy_py/src/jiminy_py/meshcat/legend.js
new file mode 100644
index 000000000..5fe1398df
--- /dev/null
+++ b/python/jiminy_py/src/jiminy_py/meshcat/legend.js
@@ -0,0 +1,79 @@
+function draggable(el) {
+  el.style.cursor = "move";
+  el.style.userSelect = "none";
+
+  el.addEventListener("mousedown", function(e) {
+    var offsetX = e.clientX - parseInt(window.getComputedStyle(this).left);
+    var offsetY = e.clientY - parseInt(window.getComputedStyle(this).top);
+
+    function mouseMoveHandler(e) {
+      el.style.top = (e.clientY - offsetY) + "px";
+      el.style.left = (e.clientX - offsetX) + "px";
+    }
+
+    function reset() {
+      window.removeEventListener("mousemove", mouseMoveHandler);
+      window.removeEventListener("mouseup", reset);
+    }
+
+    window.addEventListener("mousemove", mouseMoveHandler);
+    window.addEventListener("mouseup", reset);
+  });
+}
+
+function createLegend(id) {
+    var legend = document.createElement("div");
+    legend.id = id;
+    legend.style.position = "fixed";
+    legend.style.top = "20px";
+    legend.style.left = "20px";
+    legend.style.backgroundColor = "rgba(255, 255, 255, 0.9)";
+    legend.style.border = "2px solid";
+    legend.style.borderRadius = "5px";
+    legend.style.borderColor = "Silver";
+    legend.style.padding = "4px 10px 4px 8px";
+    document.body.prepend(legend);
+
+    draggable(legend);
+}
+
+function setLegendItem(legend, id, text, color) {
+    var box = document.createElement("div");
+    box.style.display = "inline-block";
+    box.style.height = "20px";
+    box.style.width = "20px";
+    box.style.margin = "4px 8px 4px 0px";
+    box.style.backgroundColor = color;
+
+    var label = document.createElement("span");
+    label.innerHTML = text;
+    label.style.fontFamily = "Dejavu Sans";
+    label.style.textAlign = "center";
+    label.style.alignItems = "center";
+
+    var boxContainer = document.getElementById(id);
+    if (boxContainer == null) {
+        boxContainer = document.createElement("div");
+        boxContainer.id = id;
+        boxContainer.style.display = "flex";
+        boxContainer.style.alignItems = "center";
+        legend.appendChild(boxContainer);
+    }
+    while (boxContainer.firstChild) {
+        boxContainer.removeChild(boxContainer.lastChild);
+    }
+    boxContainer.appendChild(box);
+    boxContainer.appendChild(label);
+}
+
+function removeLegendItem(legend, id) {
+    var elem = document.getElementById(id);
+    if (elem !== null) {
+        legend = elem.parentNode;
+        legend.removeChild(elem);
+        if (!legend.childElementCount)
+        {
+            document.body.removeChild(legend);
+        }
+    }
+}
diff --git a/python/jiminy_py/src/jiminy_py/meshcat/server.py b/python/jiminy_py/src/jiminy_py/meshcat/server.py
index 9f98163f1..db40f3bd3 100644
--- a/python/jiminy_py/src/jiminy_py/meshcat/server.py
+++ b/python/jiminy_py/src/jiminy_py/meshcat/server.py
@@ -12,7 +12,7 @@
 import zmq
 from zmq.eventloop.zmqstream import ZMQStream
 
-from meshcat.servers.tree import walk, find_node
+from meshcat.servers.tree import walk
 from meshcat.servers.zmqserver import (
     DEFAULT_ZMQ_METHOD, VIEWER_ROOT, StaticFileHandlerNoCache,
     ZMQWebSocketBridge, WebSocketHandler, find_available_port)
@@ -172,13 +172,6 @@ def handle_comm(self, frames: Sequence[bytes]) -> None:
                 self.is_waiting_ready_msg = False
 
     def forward_to_websockets(self, frames: Sequence[bytes]) -> None:
-        # Check if the objects are still available in cache
-        cmd, path, data = frames
-        cache_hit = (cmd == "set_object" and
-                     find_node(self.tree, path).object and
-                     find_node(self.tree, path).object == data)
-        if cache_hit:
-            return
         super().forward_to_websockets(frames)
         _, _, data = frames
         for comm_id in self.comm_pool:
diff --git a/python/jiminy_py/src/jiminy_py/meshcat/wrapper.py b/python/jiminy_py/src/jiminy_py/meshcat/wrapper.py
index 526e48142..3a78544d0 100644
--- a/python/jiminy_py/src/jiminy_py/meshcat/wrapper.py
+++ b/python/jiminy_py/src/jiminy_py/meshcat/wrapper.py
@@ -1,7 +1,10 @@
 import os
+import base64
 import atexit
 import asyncio
 import logging
+import pathlib
+import umsgpack
 import threading
 import tornado.ioloop
 from contextlib import redirect_stdout, redirect_stderr
@@ -343,6 +346,65 @@ def wait(self, require_client: bool = False) -> str:
                 process_kernel_comm()
         return self.__zmq_socket.recv().decode("utf-8")
 
+    def set_legend_item(self, uniq_id: str, color: str, text: str) -> None:
+        self.__zmq_socket.send_multipart([
+            b"set_property",      # Frontend command. Used by Python zmq server
+            b"",                  # Tree path. Empty path means root
+            umsgpack.packb({      # Backend command. Used by javascript
+                u"type": "legend",
+                u"id": uniq_id,   # Unique identifier of updated legend item
+                u"text": text,    # Any text message support by HTML5
+                u"color": color   # "rgba(0, 0, 0, 0.0)" and "black" supported
+            })
+        ])
+        self.__zmq_socket.recv()  # Receive acknowledgement
+
+    def remove_legend_item(self, uniq_id: str) -> None:
+        self.__zmq_socket.send_multipart([
+            b"set_property",
+            b"",
+            umsgpack.packb({
+                u"type": "legend",
+                u"id": uniq_id,   # Unique identifier of legend item to remove
+                u"text": ""       # Empty message means delete the item, if any
+            })
+        ])
+        self.__zmq_socket.recv()
+
+    def set_logo(self,
+                 image_fullpath: str,
+                 width: int = 100,
+                 height: int = 100) -> None:
+        file_ext = pathlib.Path(image_fullpath).suffix
+        if file_ext != ".png":
+            raise ValueError("Only .png images are supported for now.")
+
+        with open(image_fullpath, "rb") as image_file:
+            image_data = base64.b64encode(image_file.read())
+
+        self.__zmq_socket.send_multipart([
+            b"set_property",
+            b"",
+            umsgpack.packb({
+                u"type": "logo",
+                u"data": image_data.decode('utf-8'),
+                u"width": width,
+                u"height": height
+            })
+        ])
+        self.__zmq_socket.recv()
+
+    def remove_logo(self) -> None:
+        self.__zmq_socket.send_multipart([
+            b"set_property",
+            b"",
+            umsgpack.packb({
+                u"type": "logo",
+                u"data": ""   # Empty message means delete the logo, if any
+            })
+        ])
+        self.__zmq_socket.recv()
+
     def start_recording(self, fps: float, width: int, height: int) -> None:
         if not self.recorder.is_open:
             self.recorder.open()
diff --git a/python/jiminy_py/src/jiminy_py/simulator.py b/python/jiminy_py/src/jiminy_py/simulator.py
index f630f217c..9d3902655 100644
--- a/python/jiminy_py/src/jiminy_py/simulator.py
+++ b/python/jiminy_py/src/jiminy_py/simulator.py
@@ -86,7 +86,8 @@ def __init__(self,
                 "Invalid robot or controller. Make sure they are both "
                 "initialized.")
 
-        # Create shared memory for efficiency
+        # Create shared memories and python-native attribute for fast access
+        self.is_simulation_running = False
         self.stepper_state = self.engine.stepper_state
         self.system_state = self.engine.system_state
         self.sensors_data = self.robot.sensors_data
@@ -296,6 +297,9 @@ def reset(self, remove_forces: bool = False):
                               external forces. It can also be done separately.
                               Optional: Do not remove by default.
         """
+        # Consider the simulation is not longer running already
+        self.is_simulation_running = False
+
         # Reset the backend engine
         self.engine.reset(remove_forces)
 
@@ -320,10 +324,14 @@ def start(self,
                                      associated with the actual or theoretical
                                      model of the robot.
         """
+        # Call base implementation
         hresult = self.engine.start(q0, v0, None, is_state_theoretical)
         if hresult != jiminy.hresult_t.SUCCESS:
             raise RuntimeError("Failed to start the simulation.")
 
+        # Consider the simulation is now running
+        self.is_simulation_running = True
+
         # Update the observer at the end, if suitable
         if isinstance(self.controller, BaseJiminyObserverController):
             self.controller.refresh_observation(
@@ -355,6 +363,28 @@ def step(self, step_dt: float = -1) -> None:
                 self.system_state.v,
                 self.sensors_data)
 
+    def stop(self) -> None:
+        """Stop the simulation.
+
+        .. note::
+            It also releases the robot lock. Which is necessary to be able to
+            apply modifications on the robot.
+        """
+        self.engine.stop()
+        self.is_simulation_running = False
+
+    def simulate(self,
+                 t_end: float,
+                 q_init: np.ndarray,
+                 v_init: np.ndarray,
+                 a_init: Optional[np.ndarray] = None,
+                 is_state_theoretical: bool = False) -> None:
+        return_code = self.engine.simulate(
+            t_end, q_init, v_init, a_init, is_state_theoretical)
+        self.is_simulation_running = False
+        if return_code != jiminy.hresult_t.SUCCESS:
+            raise RuntimeError("The simulation failed.")
+
     def run(self,
             tf: float,
             q0: np.ndarray,
@@ -492,12 +522,13 @@ def replay(self, **kwargs: Any) -> None:
     def close(self) -> None:
         """Close the connection with the renderer.
         """
-        if hasattr(self, '_viewer') and self._viewer is not None:
-            self._viewer.close()
-            self._is_viewer_available = False
-            self._viewer = None
-        if hasattr(self, '__plot_data') and self.__plot_data is not None:
-            plt.close(self.__plot_data['fig'])
+        if hasattr(self, 'engine'):
+            if self._viewer is not None:
+                self._viewer.close()
+                self._is_viewer_available = False
+                self._viewer = None
+            if self.__plot_data is not None:
+                plt.close(self.__plot_data['fig'])
 
     def plot(self) -> None:
         """Display common simulation data over time.
diff --git a/python/jiminy_py/src/jiminy_py/viewer.py b/python/jiminy_py/src/jiminy_py/viewer.py
index 5e4aee4bb..5169dcc0c 100644
--- a/python/jiminy_py/src/jiminy_py/viewer.py
+++ b/python/jiminy_py/src/jiminy_py/viewer.py
@@ -48,6 +48,7 @@
 DEFAULT_CAPTURE_SIZE = 500
 VIDEO_FRAMERATE = 30
 VIDEO_SIZE = (1000, 1000)
+DEFAULT_LOGO_SIZE = (150, 150)
 
 DEFAULT_URDF_COLORS = {
     'green': (0.4, 0.7, 0.3, 1.0),
@@ -1327,6 +1328,9 @@ def extract_viewer_data_from_log(log_data: Dict[str, np.ndarray],
             'use_theoretical_model': use_theoretical_model}
 
 
+ColorType = Union[Tuple4FType, str]
+
+
 def play_trajectories(trajectory_data: Union[
                           TrajectoryDataType, Sequence[TrajectoryDataType]],
                       replay_speed: float = 1.0,
@@ -1340,12 +1344,14 @@ def play_trajectories(trajectory_data: Union[
                       xyz_offset: Optional[Union[
                           Tuple3FType, Sequence[Tuple3FType]]] = None,
                       urdf_rgba: Optional[Union[
-                          Tuple4FType, Sequence[Tuple4FType]]] = None,
+                          ColorType, Sequence[ColorType]]] = None,
                       backend: Optional[str] = None,
                       window_name: str = 'jiminy',
                       scene_name: str = 'world',
                       close_backend: Optional[bool] = None,
                       delete_robot_on_close: Optional[bool] = None,
+                      legend: Optional[Union[str, Sequence[str]]] = None,
+                      logo_fullpath: Optional[str] = None,
                       verbose: bool = True) -> Sequence[Viewer]:
     """Replay one or several robot trajectories in a viewer.
 
@@ -1391,7 +1397,8 @@ def play_trajectories(trajectory_data: Union[
                        Optional: None by default.
     :param urdf_rgba: List of RGBA code defining the color for each robot. It
                       will apply to every link. None to disable.
-                      Optional: Original colors of each link. No alpha.
+                      Optional: Original color if single robot, default color
+                      cycle otherwise.
     :param backend: Backend, one of 'meshcat' or 'gepetto-gui'. If None,
                     'meshcat' is used in notebook environment and 'gepetto-gui'
                     otherwise.
@@ -1408,6 +1415,16 @@ def play_trajectories(trajectory_data: Union[
     :param delete_robot_on_close: Whether or not to delete the robot from the
                                   viewer when closing it.
                                   Optional: True by default.
+    :param legend: List of text defining the legend for each robot. `urdf_rgba`
+                   must be specified to enable this option. The legend is not
+                   persistent but disabled after replay. This option is only
+                   supported by meshcat backend. None to disable.
+                   Optional: No legend if no color by default, the robots names
+                   otherwise.
+    :param logo_fullpath: Add logo/watermark to the viewer. The logo is not
+                          persistent but disabled after replay. This option is
+                          only supported by meshcat backend. None to disable.
+                          Optional: No logo by default.
     :param verbose: Add information to keep track of the process.
                     Optional: True by default.
 
@@ -1422,7 +1439,8 @@ def play_trajectories(trajectory_data: Union[
         else:
             urdf_rgba = list(islice(
                 cycle(DEFAULT_URDF_COLORS.values()), len(trajectory_data)))
-    if not isinstance(urdf_rgba, list) or isinstance(urdf_rgba[0], float):
+    if not isinstance(urdf_rgba, (list, tuple)) or \
+            isinstance(urdf_rgba[0], float):
         urdf_rgba = [urdf_rgba]
     for i, color in enumerate(urdf_rgba):
         if isinstance(color, str):
@@ -1500,6 +1518,26 @@ def play_trajectories(trajectory_data: Union[
     xyz_offset = list(xyz_offset)
     assert len(xyz_offset) == len(trajectory_data)
 
+    # Handle meshcat-specific options
+    if Viewer.backend == 'meshcat':
+        # Add legend if suitable
+        if legend is not None:
+            if not isinstance(legend, (list, tuple)):
+                legend = [legend]
+        elif urdf_rgba is not None:
+            legend = [viewer.robot_name for viewer in viewers]
+        if legend is not None:
+            assert len(legend) == len(trajectory_data)
+            for viewer, color, text in zip(viewers, urdf_rgba, legend):
+                rgba = [*[int(e * 255) for e in color[:3]], color[3]]
+                color = f"rgba({','.join(map(str, rgba))}"
+                Viewer._backend_obj.set_legend_item(
+                    viewer.robot_name, color, text)
+
+        # Add logo/watermark if requested
+        if logo_fullpath is not None:
+            Viewer._backend_obj.set_logo(logo_fullpath, *DEFAULT_LOGO_SIZE)
+
     # Do not display trajectories without data
     trajectory_data = list(trajectory_data).copy()  # No deepcopy
     for i in reversed(range(len(trajectory_data))):
@@ -1604,6 +1642,17 @@ def replay_thread(viewer, *args):
     if camera_motion is not None:
         Viewer.remove_camera_motion()
 
+    # Handle meshcat-specific options
+    if Viewer.backend == 'meshcat':
+        # Disable legend if it was enabled
+        if legend is not None:
+            for viewer in viewers:
+                Viewer._backend_obj.remove_legend_item(viewer.robot_name)
+
+        # Disable logo/watermark if it was enabled
+        if logo_fullpath is not None:
+            Viewer._backend_obj.remove_logo()
+
     # Close backend if needed
     if close_backend:
         for viewer in viewers:
diff --git a/python/jiminy_pywrap/include/jiminy/python/Jiminy.h b/python/jiminy_pywrap/include/jiminy/python/Jiminy.h
index 4d25d6589..6da01e68f 100644
--- a/python/jiminy_pywrap/include/jiminy/python/Jiminy.h
+++ b/python/jiminy_pywrap/include/jiminy/python/Jiminy.h
@@ -1677,12 +1677,9 @@ namespace python
         void visit(PyClass& cl) const
         {
             cl
-                .add_property("q", bp::make_getter(&systemState_t::q,
-                                   bp::return_internal_reference<>()))
-                .add_property("v", bp::make_getter(&systemState_t::v,
-                                   bp::return_internal_reference<>()))
-                .add_property("a", bp::make_getter(&systemState_t::a,
-                                   bp::return_internal_reference<>()))
+                .add_property("q", &PySystemStateVisitor::getPosition)
+                .add_property("v", &PySystemStateVisitor::getVelocity)
+                .add_property("a", &PySystemStateVisitor::getAcceleration)
                 .add_property("u", bp::make_getter(&systemState_t::u,
                                    bp::return_internal_reference<>()))
                 .add_property("u_motor", bp::make_getter(&systemState_t::uMotor,
@@ -1697,6 +1694,24 @@ namespace python
                 ;
         }
 
+        static bp::object getPosition(systemState_t const & self)
+        {
+            // Do not use automatic converter for efficiency
+            return convertToPython<vectorN_t>(self.q, false);
+        }
+
+        static bp::object getVelocity(systemState_t const & self)
+        {
+            // Do not use automatic converter for efficiency
+            return convertToPython<vectorN_t>(self.v, false);
+        }
+
+        static bp::object getAcceleration(systemState_t const & self)
+        {
+            // Do not use automatic converter for efficiency
+            return convertToPython<vectorN_t>(self.a, false);
+        }
+
         static std::string repr(systemState_t & self)
         {
             std::stringstream s;