Refactor UR5 environment for enhanced Nova state management and UI adjustments

- Updated the UR5 environment to include new attributes for storing Nova joint snapshots and TCP positions, improving state management.
- Modified the shutdown process to optionally preserve the Nova state, enhancing the robustness of the environment during transitions.
- Implemented methods to store and apply the last Nova state, ensuring smoother operation and better integration with user commands.
- Adjusted the UI overlay positioning in mujoco_server.py for improved visual alignment.

mujoco_server.py

@@ -1333,7 +1333,7 @@ def index():
             .overlay-tiles {
                 position: absolute;
                 right: 20px;
-                bottom: 100px;
+                bottom: 50px;
                 top: auto;
                 width: 240px;
                 display: none;

robots/ur5/ur5_env.py

@@ -154,6 +154,9 @@ class UR5Env(gym.Env):
         self._use_nova_jogging = False
         self._last_nova_sequence = None
         self._latest_nova_joints = None
+        self._nova_last_joint_snapshot = None
+        self._nova_last_tcp_position = None
+        self._nova_last_tcp_euler = None
         self._active_jog = None  # Internal jogging state
 
         self._nova_stream_interval_sec = 0.2
@@ -167,6 +170,8 @@ class UR5Env(gym.Env):
 
         # Remember the requested Nova configuration so we can toggle it later
         self._nova_api_config = dict(nova_api_config) if nova_api_config else None
+        self._target_pos = None
+        self._target_euler = None
         self._nova_user_enabled = False
         if self._nova_api_config:
             enabled = self.set_nova_enabled(True)
@@ -301,7 +306,7 @@ class UR5Env(gym.Env):
             self._shutdown_nova()
             return False
 
-    def _shutdown_nova(self) -> None:
+    def _shutdown_nova(self, preserve_snapshot: bool = False) -> None:
         """Stop any active Nova clients and reset Nova-dependent flags."""
         if self._nova_jogger is not None:
             try:
@@ -323,6 +328,35 @@ class UR5Env(gym.Env):
         self._latest_nova_joints = None
         self._last_nova_sequence = None
         self._nova_user_enabled = False
+        if not preserve_snapshot:
+            self._nova_last_joint_snapshot = None
+            self._nova_last_tcp_position = None
+            self._nova_last_tcp_euler = None
+
+    def _store_nova_snapshot(self) -> None:
+        if self._latest_nova_joints is not None:
+            self._nova_last_joint_snapshot = self._latest_nova_joints.copy()
+        if self._target_pos is not None:
+            self._nova_last_tcp_position = self._target_pos.copy()
+        if self._target_euler is not None:
+            self._nova_last_tcp_euler = self._target_euler.copy()
+
+    def _apply_last_nova_state_to_internal(self) -> None:
+        applied = False
+        if self._nova_last_joint_snapshot is not None:
+            self.data.qpos[:6] = self._nova_last_joint_snapshot.copy()
+            self.data.ctrl[:6] = self._nova_last_joint_snapshot.copy()
+            self._joint_targets = self._nova_last_joint_snapshot.copy()
+            applied = True
+        if self._nova_last_tcp_position is not None:
+            self._target_pos = self._nova_last_tcp_position.copy()
+            self.data.mocap_pos[0] = self._target_pos
+            applied = True
+        if self._nova_last_tcp_euler is not None:
+            self._target_euler = self._nova_last_tcp_euler.copy()
+            applied = True
+        if applied:
+            mujoco.mj_forward(self.model, self.data)
 
     def set_nova_enabled(self, enabled: bool) -> bool:
         """Enable or disable Nova integration according to the stored configuration."""
@@ -331,11 +365,14 @@ class UR5Env(gym.Env):
                 print("[Nova] Cannot enable Nova API: configuration unavailable.")
                 self._nova_user_enabled = False
                 return False
+            self._shutdown_nova()
             success = self._configure_nova_api(self._nova_api_config)
             self._nova_user_enabled = success
             return success
 
-        self._shutdown_nova()
+        self._store_nova_snapshot()
+        self._shutdown_nova(preserve_snapshot=True)
+        self._apply_last_nova_state_to_internal()
         self._nova_user_enabled = False
         return True
 
@@ -966,5 +1003,6 @@ class UR5Env(gym.Env):
                         )
                         self._target_euler = np.array([roll, pitch, yaw], dtype=np.float32)
 
-        self._update_nova_interpolation(time.time())
+            self._update_nova_interpolation(time.time())
+        self._store_nova_snapshot()
         return self._latest_nova_joints is not None