Add environment loading utility and integrate into Nova API and Jogger

- Introduced a new utility module for loading environment variables from `.env.local` to streamline configuration management.
- Integrated the environment loading functionality into the Nova API and Jogger modules to ensure proper configuration during runtime.
- Updated the UR5 environment to utilize the loaded home pose for joint initialization, enhancing simulation consistency.

robots/ur5/env_loader.py

@@ -0,0 +1,22 @@
+"""
+Utility helpers for loading Nova-related environment files.
+"""
+
+import os
+from pathlib import Path
+
+
+def load_local_env() -> None:
+    """Load `.env.local` from the nova-sim root so configs pick up saved credentials."""
+    root = Path(__file__).resolve().parents[2]
+    env_path = root / ".env.local"
+    if not env_path.exists():
+        return
+
+    with env_path.open("r") as fp:
+        for raw in fp:
+            line = raw.strip()
+            if not line or line.startswith("#") or "=" not in line:
+                continue
+            key, value = line.split("=", 1)
+            os.environ.setdefault(key.strip(), value.strip())

robots/ur5/nova_api.py

@@ -1,3 +1,4 @@
+import importlib.util
 import json
 import math
 import os
@@ -6,6 +7,7 @@ import time
 import urllib.error
 import urllib.request
 from dataclasses import dataclass
+from pathlib import Path
 from typing import Any, Dict, Optional, Tuple
 
 try:
@@ -16,6 +18,21 @@ except Exception as e:  # pragma: no cover - optional dependency
     connect = None
 
 
+def _load_local_env_module():
+    try:
+        from robots.ur5 import env_loader
+    except ImportError:
+        env_path = Path(__file__).resolve().parent / "env_loader.py"
+        spec = importlib.util.spec_from_file_location("robots.ur5.env_loader", env_path)
+        module = importlib.util.module_from_spec(spec)
+        spec.loader.exec_module(module)
+        env_loader = module
+    env_loader.load_local_env()
+
+
+_load_local_env_module()
+
+
 @dataclass
 class NovaApiConfig:
     instance_url: str
@@ -103,6 +120,7 @@ class NovaApiClient:
         self._state_thread: Optional[threading.Thread] = None
         self._state_ws = None
         self._motion_group_description: Optional[Dict[str, Any]] = None
+        self._last_state_log_time = 0.0
 
     def start_state_stream(self) -> None:
         if connect is None:
@@ -187,6 +205,20 @@ class NovaApiClient:
                         continue
                     with self._state_lock:
                         self._latest_state = data
+                        now = time.time()
+                        if now - self._last_state_log_time >= 10.0:
+                            self._last_state_log_time = now
+                            result = data.get("result", data)
+                            seq = result.get("sequence_number", "N/A")
+                            joints = result.get("joint_position")
+                            joint_info = f"{len(joints)} joints" if isinstance(joints, list) else "no joints"
+                            payload = {
+                                "seq": seq,
+                                "joint_position": joints,
+                                "timestamp": result.get("timestamp")
+                            }
+                            payload_str = json.dumps(payload, default=str)
+                            print(f"[Nova] State stream update (seq={seq}, {joint_info}) payload={payload_str}")
                 ws.close()
             except Exception as exc:
                 print(f"Nova state stream error: {exc}")

robots/ur5/nova_jogger.py

@@ -3,11 +3,13 @@ Nova Jogger Client
 Handles bidirectional jogging control for Nova API motion groups.
 """
 
+import importlib.util
 import json
 import threading
 import time
 import urllib.error
 import urllib.request
+from pathlib import Path
 from typing import Optional, Literal, Dict, Any
 
 try:
@@ -17,6 +19,21 @@ except Exception as e:
     connect = None
 
 
+def _load_local_env_module():
+    try:
+        from robots.ur5 import env_loader
+    except ImportError:
+        env_path = Path(__file__).resolve().parent / "env_loader.py"
+        spec = importlib.util.spec_from_file_location("robots.ur5.env_loader", env_path)
+        module = importlib.util.module_from_spec(spec)
+        spec.loader.exec_module(module)
+        env_loader = module
+    env_loader.load_local_env()
+
+
+_load_local_env_module()
+
+
 JoggingMode = Literal["joint", "cartesian"]
 JogDirection = Literal["positive", "negative", "+", "-"]  # Accept both formats, normalize to positive/negative
 CartesianAxis = Literal["x", "y", "z"]

robots/ur5/ur5_env.py

@@ -8,6 +8,7 @@ import mujoco
 import os
 import importlib.util
 import math
+import xml.etree.ElementTree as ET
 import random
 
 # Import IK controller
@@ -32,7 +33,7 @@ TABLE_PAD_TOP = 0.35  # meters: upper boundary (positive Y) for randomized targe
 TABLE_SURFACE_Z = 0.53  # meters: actual table-top surface height from the XML scene
 TABLE_OBJECT_HEIGHT = 0.03  # meters: height of the movable T-object along Z (from scene)
 TABLE_MIN_DISTANCE = 0.12  # meters: minimum planar distance between target and object at reset
-TARGET_HEIGHT_OFFSET = 0.3  # meters: raise IK target above the goal/object when initializing
+TARGET_HEIGHT_OFFSET = 0.05  # raise IK target slightly above the goal/object so the stick ends close to the table
 TABLE_OBJECT_SPAWN_Z = TABLE_SURFACE_Z + TABLE_OBJECT_HEIGHT  # center height so object rests on table
 
 
@@ -163,8 +164,10 @@ class UR5Env(gym.Env):
         )
 
         # Set initial joint positions to home pose and compute FK
-        self.data.qpos[:6] = self.DEFAULT_HOME_POSE.copy()
-        self.data.ctrl[:6] = self.DEFAULT_HOME_POSE.copy()
+        loaded_pose = self._load_home_pose(model_path)
+        self.home_pose = loaded_pose if loaded_pose is not None else self.DEFAULT_HOME_POSE.copy()
+        self.data.qpos[:6] = self.home_pose.copy()
+        self.data.ctrl[:6] = self.home_pose.copy()
         mujoco.mj_forward(self.model, self.data)
 
         # Target position/orientation from FK of home pose
@@ -185,7 +188,7 @@ class UR5Env(gym.Env):
         # Control mode: 'ik' (end-effector target) or 'joint' (direct joint positions)
         self._control_mode = 'ik'
         # Direct joint targets (used when control_mode is 'joint')
-        self._joint_targets = self.DEFAULT_HOME_POSE.copy()
+        self._joint_targets = self._compute_joint_targets()
 
         # Configure Nova API if requested (must be after Nova attributes are initialized)
         if nova_api_config:
@@ -287,6 +290,25 @@ class UR5Env(gym.Env):
         self.data.qpos[qpos_adr:qpos_adr + 3] = pos
         self.data.qpos[qpos_adr + 3:qpos_adr + 7] = quat
 
+    def _load_home_pose(self, xml_path: str) -> np.ndarray | None:
+        """Load the UR5 home qpos from the scene XML keyframe named 'home'."""
+        try:
+            tree = ET.parse(xml_path)
+            root = tree.getroot()
+        except Exception:
+            return None
+        for key in root.findall(".//key"):
+            if key.get("name") != "home":
+                continue
+            qpos = key.get("qpos")
+            if not qpos:
+                continue
+            values = [float(val) for val in qpos.split()]
+            if len(values) < 6:
+                continue
+            return np.array(values[:6], dtype=np.float32)
+        return None
+
     @staticmethod
     def _sample_table_position():
         x = random.uniform(TABLE_PAD_LEFT, TABLE_PAD_RIGHT)
@@ -325,6 +347,9 @@ class UR5Env(gym.Env):
             x, y, z: Target position in meters
             update_joint_targets: If True, compute IK and update joint targets
         """
+        if self._use_nova_state_stream and self._nova_client is not None:
+            print("[Nova] Ignoring local target updates while Nova state stream is active")
+            return
         self._target_pos = np.array([x, y, z], dtype=np.float32)
         # Update mocap body position for visualization
         self.data.mocap_pos[0] = self._target_pos
@@ -397,7 +422,7 @@ class UR5Env(gym.Env):
             try:
                 if jog_type == 'joint':
                     return self._nova_jogger.start_joint_jog(
-                        joint=kwargs.get('joint', 0),
+                        joint=max(0, min(5, kwargs.get('joint', 1) - 1)),
                         direction=kwargs.get('direction', '+'),
                         velocity_rads_per_sec=kwargs.get('velocity', 0.5)
                     )
@@ -613,8 +638,8 @@ class UR5Env(gym.Env):
         mujoco.mj_resetData(self.model, self.data)
 
         # Set to home pose
-        self.data.qpos[:6] = self.DEFAULT_HOME_POSE.copy()
-        self.data.ctrl[:6] = self.DEFAULT_HOME_POSE.copy()
+        self.data.qpos[:6] = self.home_pose.copy()
+        self.data.ctrl[:6] = self.home_pose.copy()
         if self.has_gripper:
             self.data.ctrl[6] = 0  # Gripper open (Robotiq: 0=open, 255=closed)
 
@@ -645,7 +670,7 @@ class UR5Env(gym.Env):
         self.data.mocap_pos[0] = self._target_pos
         if self.has_gripper:
             self._gripper_target = 0.0  # Open gripper
-        self._joint_targets = self.DEFAULT_HOME_POSE.copy()
+        self._joint_targets = self.home_pose.copy()
 
         self.steps = 0
 
@@ -654,8 +679,19 @@ class UR5Env(gym.Env):
 
     def step(self, action):
         """Step with explicit action (for RL training)."""
-        # Apply action to actuators
-        self.data.ctrl[:] = action
+        if self._use_nova_state_stream and self._nova_client is not None:
+            synced = self._sync_nova_state()
+            if synced and self._latest_nova_joints is not None:
+                self.data.ctrl[:6] = self._latest_nova_joints
+                # Keep the joints aligned with Nova even though the RL action is ignored
+                self.data.qpos[:6] = self._latest_nova_joints
+                self.data.qvel[:6] = 0.0
+            else:
+                # Nova data unavailable; fall back to the provided action
+                self.data.ctrl[:] = action
+        else:
+            # Apply action to actuators (default RL mode)
+            self.data.ctrl[:] = action
 
         # Step simulation
         mujoco.mj_step(self.model, self.data)
@@ -802,5 +838,19 @@ class UR5Env(gym.Env):
         self.data.qvel[:6] = 0.0
         self.data.ctrl[:6] = self.data.qpos[:6]
         mujoco.mj_forward(self.model, self.data)
+        tcp_pose = result.get("tcp_pose") or {}
+        pos = tcp_pose.get("position")
+        if pos and len(pos) == 3:
+            self._target_pos = np.array(pos, dtype=np.float32)
+            self.data.mocap_pos[0] = self._target_pos
+        orientation = tcp_pose.get("orientation")
+        if orientation and len(orientation) == 4:
+            roll, pitch, yaw = IKController.quat_to_euler(
+                float(orientation[0]),
+                float(orientation[1]),
+                float(orientation[2]),
+                float(orientation[3]),
+            )
+            self._target_euler = np.array([roll, pitch, yaw], dtype=np.float32)
         self._last_nova_sequence = sequence
         return True

tests/test_nova_functionality.py

@@ -0,0 +1,99 @@
+import os
+import time
+
+import pytest
+from pathlib import Path
+
+from robots.ur5 import nova_api, nova_jogger
+
+
+def _load_env():
+    """Load `.env.local` if present so the Nova config can pull its values."""
+    env_path = Path(__file__).parent.parent / ".env.local"
+    if not env_path.exists():
+        return
+    with env_path.open("r") as fp:
+        for raw in fp:
+            line = raw.strip()
+            if not line or line.startswith("#") or "=" not in line:
+                continue
+            key, value = line.split("=", 1)
+            os.environ.setdefault(key.strip(), value.strip())
+
+
+_load_env()
+
+
+@pytest.fixture(scope="module")
+def nova_config():
+    """Provide NovaApiConfig loaded from environment or skip if not configured."""
+    try:
+        return nova_api.NovaApiConfig.from_env()
+    except ValueError as exc:
+        pytest.skip(f"Nova configuration missing: {exc}")
+
+
+def _create_jogger(config: nova_api.NovaApiConfig) -> nova_jogger.NovaJogger:
+    return nova_jogger.NovaJogger(
+        config.instance_url,
+        config.access_token,
+        config.cell_id,
+        config.controller_id,
+        config.motion_group_id,
+    )
+
+
+def test_state_stream_receives_updates(nova_config):
+    """Ensure the state stream connects and delivers joint positions."""
+    if nova_api.connect is None:
+        pytest.skip("websockets package required for Nova state stream")
+
+    client = nova_api.NovaApiClient(nova_config)
+    client.start_state_stream()
+    try:
+        deadline = time.time() + 10.0
+        state = None
+        while time.time() < deadline:
+            state = client.get_latest_state()
+            if state:
+                break
+            time.sleep(0.1)
+        assert state, "No state updates were received"
+        result = state.get("result", state)
+        assert "joint_position" in result, "Expected joint_position field in state payload"
+        assert isinstance(result["joint_position"], list)
+        assert client.is_state_stream_connected()
+    finally:
+        client.stop()
+
+
+def test_jogger_joint_connection(nova_config):
+    """Connect the jogger and exercise a zero-velocity joint command."""
+    if nova_jogger.connect is None:
+        pytest.skip("websockets package required for Nova jogging")
+
+    jogger = _create_jogger(nova_config)
+    try:
+        assert jogger.connect()
+        assert jogger.start_joint_jog(0, "+", velocity_rads_per_sec=0.0)
+        assert jogger.is_connected()
+        time.sleep(0.2)
+        assert jogger.stop()
+    finally:
+        jogger.disconnect()
+
+
+def test_jogger_cartesian_connection(nova_config):
+    """Connect the jogger and exercise a zero-velocity cartesian command."""
+    if nova_jogger.connect is None:
+        pytest.skip("websockets package required for Nova jogging")
+
+    jogger = _create_jogger(nova_config)
+    try:
+        assert jogger.connect()
+        assert jogger.start_cartesian_translation("z", "+", velocity_mm_per_sec=0.0)
+        assert jogger.is_connected()
+        time.sleep(0.2)
+        assert jogger.stop()
+    finally:
+        jogger.disconnect()

tests/test_nova_state_stream_commands.py

@@ -0,0 +1,124 @@
+import math
+import time
+
+import pytest
+
+from robots.ur5 import env_loader, nova_api, nova_jogger
+
+
+env_loader.load_local_env()
+
+
+def _get_result(state: dict) -> dict:
+    return state.get("result", state)
+
+
+def _wait_for_state_condition(client: nova_api.NovaApiClient, predicate, timeout: float = 10.0) -> dict:
+    deadline = time.monotonic() + timeout
+    last_seq = None
+    while time.monotonic() < deadline:
+        state = client.get_latest_state()
+        if not state:
+            time.sleep(0.1)
+            continue
+        result = _get_result(state)
+        seq = result.get("sequence_number")
+        if seq is not None and seq == last_seq:
+            time.sleep(0.1)
+            continue
+        last_seq = seq
+        if predicate(result):
+            return result
+        time.sleep(0.1)
+    pytest.fail("Timed out waiting for the requested Nova state")
+
+
+def _joints_changed(result: dict, baseline: list[float], threshold: float = 0.01) -> bool:
+    joints = result.get("joint_position") or []
+    if len(joints) < len(baseline):
+        return False
+    return any(abs(joints[i] - baseline[i]) >= threshold for i in range(len(baseline)))
+
+
+def _tcp_moved(result: dict, baseline_pos: list[float], threshold: float = 0.01) -> bool:
+    tcp_pose = result.get("tcp_pose") or {}
+    pos = tcp_pose.get("position")
+    if not pos or len(pos) < 3:
+        return False
+    return math.dist(pos[:3], baseline_pos[:3]) >= threshold
+
+
+@pytest.fixture(scope="module")
+def nova_config():
+    try:
+        return nova_api.NovaApiConfig.from_env()
+    except ValueError as exc:
+        pytest.skip(f"Nova configuration missing: {exc}")
+
+
+def _create_jogger(config: nova_api.NovaApiConfig) -> nova_jogger.NovaJogger:
+    jogger = nova_jogger.NovaJogger(
+        config.instance_url,
+        config.access_token,
+        config.cell_id,
+        config.controller_id,
+        config.motion_group_id,
+    )
+    assert jogger.connect()
+    return jogger
+
+
+def test_state_stream_updates_after_joint_jog(nova_config):
+    if nova_api.connect is None or nova_jogger.connect is None:
+        pytest.skip("websockets dependency required for state stream and jogging")
+
+    client = nova_api.NovaApiClient(nova_config)
+    client.start_state_stream()
+    jogger = _create_jogger(nova_config)
+    try:
+        baseline = _wait_for_state_condition(client, lambda _: True, timeout=15.0)
+        baseline_joints = baseline.get("joint_position") or []
+        assert len(baseline_joints) >= 6, "State stream must report 6 joint positions"
+
+        assert jogger.start_joint_jog(0, "+", velocity_rads_per_sec=0.2)
+        time.sleep(1.0)
+        assert jogger.stop()
+
+        updated = _wait_for_state_condition(
+            client,
+            lambda result: _joints_changed(result, baseline_joints, threshold=0.008),
+            timeout=15.0,
+        )
+        assert _joints_changed(updated, baseline_joints, threshold=0.008)
+    finally:
+        jogger.disconnect()
+        client.stop()
+
+
+def test_state_stream_reflects_tcp_translation(nova_config):
+    if nova_api.connect is None or nova_jogger.connect is None:
+        pytest.skip("websockets dependency required for state stream and jogging")
+
+    client = nova_api.NovaApiClient(nova_config)
+    client.start_state_stream()
+    jogger = _create_jogger(nova_config)
+    try:
+        baseline = _wait_for_state_condition(client, lambda state: bool(state.get("tcp_pose")), timeout=15.0)
+        tcp_pose = baseline.get("tcp_pose") or {}
+        position = tcp_pose.get("position")
+        if not position or len(position) < 3:
+            pytest.skip("TCP pose not available from state stream")
+
+        assert jogger.start_cartesian_translation("z", "+", velocity_mm_per_sec=20.0)
+        time.sleep(1.2)
+        assert jogger.stop()
+
+        moved = _wait_for_state_condition(
+            client,
+            lambda result: _tcp_moved(result, position, threshold=0.005),
+            timeout=15.0,
+        )
+        assert _tcp_moved(moved, position, threshold=0.005)
+    finally:
+        jogger.disconnect()
+        client.stop()

tests/test_ur5_home_pose.py

@@ -0,0 +1,18 @@
+import mujoco
+import numpy as np
+from robots.ur5.ur5_env import UR5Env
+
+
+def test_ur5_ik_targets_match_home_pose():
+    """Ensure the IK solution satisfies the Cartesian pose derived from the XML home qpos."""
+    env = UR5Env(render_mode=None, width=1, height=1, scene_name="scene_t_push")
+    try:
+        ik_solution = env._compute_joint_targets()
+
+        env.data.qpos[:6] = ik_solution
+        mujoco.mj_forward(env.model, env.data)
+
+        ee_pos = env.get_end_effector_pos()
+        assert np.allclose(ee_pos, env._target_pos, atol=1e-2)
+    finally:
+        env.close()