feat(go2): rage mode toggle + lowstate stream & battery SOC

[ruthwikdasyam]

Problem

Three gaps in the Go2 connection layer:

• enable_rage_mode() could only turn Rage Mode on — no way back to the normal velocity envelope.
• lowstate_stream() (battery, IMU, etc.) exists on the WebRTC connection but isn't part of Go2ConnectionProtocol, so callers can't use it type-safely and the sim backends don't implement it.
• No way to read the robot's battery level.

Solution

• Rage toggle: renamed enable_rage_mode() → set_rage_mode(enable: bool), bidirectional across the WebRTC connection, protocol, RPC wrapper, and sim backends (mujoco, dimsim, replay). On disable it sends 2059 {data:False} + SwitchJoystick(False) to restore the normal envelope; BalanceStand precondition is re-established inside the WebRTC connection.
• lowstate_stream: added to Go2ConnectionProtocol and satisfied on the sims (mujoco → empty(), dimsim → Subject(), replay → empty()). Real WebRTC impl already exists on main and is unchanged.
• Battery SOC: GO2Connection subscribes to the lowstate stream, caches bms_state.soc, and exposes it as a get_battery_soc() skill (0-100, or None until first lowstate). Degrades cleanly to None on sim/replay.

How to Test

dimos run unitree-go2-webrtc-rage-keyboard-teleop

Vertified in upcoming hosted_teleop PR - accessing battery info from low state, and turning ON/OFF rage mode

Contributor License Agreement

• I have read and approved the CLA.

[codecov]

❌ 1 Tests Failed:

Tests completed	Failed	Passed	Skipped
2335	1	2334	74

View the top 1 failed test(s) by shortest run time

dimos.control.test_control.TestIntegration::test_full_trajectory_execution

Stack Traces | 0.793s run time

self = <dimos.control.test_control.TestIntegration object at 0x7f97fabcfb30>
mock_adapter = <MagicMock spec='ManipulatorAdapter' id='140288324374928'>

    def test_full_trajectory_execution(self, mock_adapter):
        component = HardwareComponent(
            hardware_id="arm",
            hardware_type=HardwareType.MANIPULATOR,
            joints=make_joints("arm", 6),
        )
        hw = ConnectedHardware(mock_adapter, component)
        hardware = {"arm": hw}
    
        config = JointTrajectoryTaskConfig(
            joint_names=[f"arm/joint{i + 1}" for i in range(6)],
            priority=10,
        )
        traj_task = JointTrajectoryTask(name="traj_arm", config=config)
        tasks = {"traj_arm": traj_task}
    
        joint_to_hardware = {f"arm/joint{i + 1}": "arm" for i in range(6)}
    
        tick_loop = TickLoop(
            tick_rate=100.0,
            hardware=hardware,
            hardware_lock=threading.Lock(),
            tasks=tasks,
            task_lock=threading.Lock(),
            joint_to_hardware=joint_to_hardware,
        )
    
        trajectory = JointTrajectory(
            joint_names=[f"arm/joint{i + 1}" for i in range(6)],
            points=[
                TrajectoryPoint(
                    positions=[0.0] * 6,
                    velocities=[0.0] * 6,
                    time_from_start=0.0,
                ),
                TrajectoryPoint(
                    positions=[0.5] * 6,
                    velocities=[0.0] * 6,
                    time_from_start=0.5,
                ),
            ],
        )
    
        tick_loop.start()
        traj_task.execute(trajectory)
    
        time.sleep(0.6)
        tick_loop.stop()
    
>       assert traj_task.get_state() == TrajectoryState.COMPLETED
E       assert <TrajectoryState.EXECUTING: 1> == <TrajectoryState.COMPLETED: 2>
E        +  where <TrajectoryState.EXECUTING: 1> = get_state()
E        +    where get_state = <dimos.control.tasks.trajectory_task.trajectory_task.JointTrajectoryTask object at 0x7f976bb97dd0>.get_state
E        +  and   <TrajectoryState.COMPLETED: 2> = TrajectoryState.COMPLETED

component  = HardwareComponent(hardware_id='arm', hardware_type=<HardwareType.MANIPULATOR: 'manipulator'>, joints=['arm/joint1', 'a...adapter_type='mock', address=None, auto_enable=True, gripper_joints=[], domain_id=0, adapter_kwargs={}, wb_config=None)
config     = JointTrajectoryTaskConfig(joint_names=['arm/joint1', 'arm/joint2', 'arm/joint3', 'arm/joint4', 'arm/joint5', 'arm/joint6'], priority=10)
hardware   = {'arm': <dimos.control.hardware_interface.ConnectedHardware object at 0x7f976bb97da0>}
hw         = <dimos.control.hardware_interface.ConnectedHardware object at 0x7f976bb97da0>
joint_to_hardware = {'arm/joint1': 'arm', 'arm/joint2': 'arm', 'arm/joint3': 'arm', 'arm/joint4': 'arm', ...}
mock_adapter = <MagicMock spec='ManipulatorAdapter' id='140288324374928'>
self       = <dimos.control.test_control.TestIntegration object at 0x7f97fabcfb30>
tasks      = {'traj_arm': <dimos.control.tasks.trajectory_task.trajectory_task.JointTrajectoryTask object at 0x7f976bb97dd0>}
tick_loop  = <dimos.control.tick_loop.TickLoop object at 0x7f976bb97e90>
traj_task  = <dimos.control.tasks.trajectory_task.trajectory_task.JointTrajectoryTask object at 0x7f976bb97dd0>
trajectory = JointTrajectory(points=[TrajectoryPoint(time_from_start=0.0, positions=[0.0, 0.0, 0.0, 0.0, 0.0, 0.0], velocities=[0.0...mes=['arm/joint1', 'arm/joint2', 'arm/joint3', 'arm/joint4', 'arm/joint5', 'arm/joint6'], timestamp=1782425436.0188289)

dimos/control/test_control.py:584: AssertionError

View the full list of 1 ❄️ flaky test(s)

dimos.e2e_tests.test_dimsim_spatial_memory::test_go_to_the_bed

Flake rate in main: 23.40% (Passed 36 times, Failed 11 times)

Stack Traces | 569s run time

lcm_spy = <dimos.e2e_tests.lcm_spy.LcmSpy object at 0x789d4fc8f2c0>
start_blueprint = <function start_blueprint.<locals>.set_name_and_start at 0x789d4d32d800>
human_input = <function human_input.<locals>.send_human_input at 0x789d4d32d8a0>
dim_sim = <dimos.e2e_tests.dim_sim_client.DimSimClient object at 0x789d4f12e8d0>
explore_house = <function explore_house.<locals>.explore at 0x789d4d32dda0>

    @pytest.mark.self_hosted_large
    def test_go_to_the_bed(lcm_spy, start_blueprint, human_input, dim_sim, explore_house) -> None:
        start_blueprint(
            "run",
            "unitree-go2-agentic",
            simulator="dimsim",
        )
        lcm_spy.save_topic(".../McpClient/on_system_modules/res")
        lcm_spy.wait_for_saved_topic(".../McpClient/on_system_modules/res", timeout=1200.0)
    
        explore_house()
    
        human_input("go to the bed")
    
>       lcm_spy.wait_until_odom_position(-3.567, -1.332, threshold=2, timeout=180)

dim_sim    = <dimos.e2e_tests.dim_sim_client.DimSimClient object at 0x789d4f12e8d0>
explore_house = <function explore_house.<locals>.explore at 0x789d4d32dda0>
human_input = <function human_input.<locals>.send_human_input at 0x789d4d32d8a0>
lcm_spy    = <dimos.e2e_tests.lcm_spy.LcmSpy object at 0x789d4fc8f2c0>
start_blueprint = <function start_blueprint.<locals>.set_name_and_start at 0x789d4d32d800>

dimos/e2e_tests/test_dimsim_spatial_memory.py:32: 
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 
dimos/e2e_tests/lcm_spy.py:182: in wait_until_odom_position
    self.wait_for_message_result(
        predicate  = <function LcmSpy.wait_until_odom_position.<locals>.predicate at 0x789d4d32df80>
        self       = <dimos.e2e_tests.lcm_spy.LcmSpy object at 0x789d4fc8f2c0>
        threshold  = 2
        timeout    = 180
        x          = -3.567
        y          = -1.332
dimos/e2e_tests/lcm_spy.py:168: in wait_for_message_result
    self.wait_until(
        event      = <threading.Event at 0x789d4f12fe60: unset>
        fail_message = 'Failed to get to position x=-3.567, y=-1.332'
        listener   = <function LcmSpy.wait_for_message_result.<locals>.listener at 0x789d4d32de40>
        predicate  = <function LcmSpy.wait_until_odom_position.<locals>.predicate at 0x789d4d32df80>
        self       = <dimos.e2e_tests.lcm_spy.LcmSpy object at 0x789d4fc8f2c0>
        timeout    = 180
        topic      = '/odom#geometry_msgs.PoseStamped'
        type       = <class 'dimos.msgs.geometry_msgs.PoseStamped.PoseStamped'>
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 

self = <dimos.e2e_tests.lcm_spy.LcmSpy object at 0x789d4fc8f2c0>

    def wait_until(
        self,
        *,
        condition: Callable[[], bool],
        timeout: float,
        error_message: str,
        poll_interval: float = 0.1,
    ) -> None:
        start_time = time.time()
        while time.time() - start_time < timeout:
            if condition():
                return
            time.sleep(poll_interval)
>       raise TimeoutError(error_message)
E       TimeoutError: Failed to get to position x=-3.567, y=-1.332

condition  = <bound method Event.is_set of <threading.Event at 0x789d4f12fe60: unset>>
error_message = 'Failed to get to position x=-3.567, y=-1.332'
poll_interval = 0.1
self       = <dimos.e2e_tests.lcm_spy.LcmSpy object at 0x789d4fc8f2c0>
start_time = 1782430658.138446
timeout    = 180

dimos/e2e_tests/lcm_spy.py:105: TimeoutError

To view more test analytics, go to the Test Analytics Dashboard
_{📋 Got 3 mins? Take this short survey to help us improve Test Analytics.}

[greptile-apps]

Greptile Summary

This PR completes three gaps in the Go2 connection layer: renames enable_rage_mode() to set_rage_mode(enable: bool) to support bidirectional toggling, adds lowstate_stream() to Go2ConnectionProtocol with appropriate sim stubs, and exposes a new get_battery_soc() skill that caches the latest BMS SOC from the lowstate stream.

• Rage toggle: set_rage_mode is now bidirectional across the WebRTC connection, protocol, RPC wrapper, and all sim backends; the balance_stand() precondition and settle delays are handled inside UnitreeWebRTCConnection.set_rage_mode.
• lowstate_stream protocol extension: Go2ConnectionProtocol now includes lowstate_stream() -> Observable[LowStateMsg]; MujocoConnection and ReplayConnection return empty(), while DimSimConnection returns a live Subject() consistent with its other streams.
• Battery SOC skill: GO2Connection subscribes to the lowstate stream on start(), caches each message in _latest_lowstate, and exposes it via a @skill-decorated get_battery_soc() with graceful None fallback.

Confidence Score: 5/5

Safe to merge — the changes are additive and well-scoped, with correct fallbacks on all sim backends.

The rename from enable_rage_mode to set_rage_mode(enable) is consistently applied across all four files. The lowstate_stream addition is correctly stubbed on every backend, and get_battery_soc has a broad try/except that degrades gracefully to None. No regressions are introduced in the existing WebRTC path.

No files require special attention.

Important Files Changed

Filename	Overview
dimos/robot/unitree/go2/connection.py	Core changes: adds lowstate_stream to Go2ConnectionProtocol, renames enable_rage_mode → set_rage_mode(enable), wires lowstate_stream subscription in start(), and adds get_battery_soc() skill with try/except guard. All sim backends and the protocol are kept in sync.
dimos/robot/unitree/connection.py	Replaces enable_rage_mode() with set_rage_mode(enable: bool) on the WebRTC layer; adds balance_stand() precondition with warning log, parameterises both the 2059 API call and SwitchJoystick, and skips the 2-second settle delay when disabling.
dimos/robot/unitree/mujoco_connection.py	Renames enable_rage_mode and adds lowstate_stream() returning empty() under @functools.cache, consistent with the existing stream pattern in this class.
dimos/robot/unitree/dimsim_connection.py	Renames enable_rage_mode and adds lowstate_stream() returning Subject() under @functools.cache, consistent with DimSim's live-data Subject pattern for other streams.

Sequence Diagram

%%{init: {'theme': 'neutral'}}%%
sequenceDiagram
    participant Caller
    participant GO2Connection
    participant WebRTC as UnitreeWebRTCConnection
    participant Robot

    Note over GO2Connection,Robot: start() — lowstate subscription
    GO2Connection->>WebRTC: lowstate_stream()
    WebRTC-->>GO2Connection: Observable[LowStateMsg]
    Robot-->>GO2Connection: _on_lowstate(msg) [background]
    GO2Connection->>GO2Connection: "_latest_lowstate = msg"

    Note over Caller,Robot: get_battery_soc()
    Caller->>GO2Connection: get_battery_soc()
    GO2Connection->>GO2Connection: _latest_lowstate["data"]["bms_state"]["soc"]
    GO2Connection-->>Caller: int (0-100) or None

    Note over Caller,Robot: set_rage_mode(enable)
    Caller->>GO2Connection: set_rage_mode(enable)
    GO2Connection->>WebRTC: set_rage_mode(enable)
    WebRTC->>Robot: balance_stand()
    WebRTC->>WebRTC: sleep(0.3s)
    WebRTC->>Robot: "2059 {data: enable}"
    alt "enable == True"
        WebRTC->>WebRTC: sleep(2.0s)
    end
    WebRTC->>Robot: SwitchJoystick(enable)
    WebRTC-->>GO2Connection: bool
    GO2Connection-->>Caller: bool

Loading

%%{init: {'theme': 'base', 'themeVariables': {"darkMode": true, "background": "#0d1117", "primaryColor": "#21262d", "primaryTextColor": "#e6edf3", "primaryBorderColor": "#8b949e", "lineColor": "#8b949e", "textColor": "#e6edf3", "edgeLabelBackground": "#161b22", "actorBkg": "#21262d", "actorBorder": "#8b949e", "actorTextColor": "#e6edf3", "actorLineColor": "#8b949e", "signalColor": "#8b949e", "signalTextColor": "#e6edf3", "noteBkgColor": "#373320", "noteBorderColor": "#d4a72c", "noteTextColor": "#f0e6c0", "labelBoxBkgColor": "#21262d", "labelBoxBorderColor": "#8b949e", "labelTextColor": "#e6edf3", "loopTextColor": "#e6edf3", "activationBkgColor": "#30363d", "activationBorderColor": "#8b949e"}}}%%
sequenceDiagram
    participant Caller
    participant GO2Connection
    participant WebRTC as UnitreeWebRTCConnection
    participant Robot

    Note over GO2Connection,Robot: start() — lowstate subscription
    GO2Connection->>WebRTC: lowstate_stream()
    WebRTC-->>GO2Connection: Observable[LowStateMsg]
    Robot-->>GO2Connection: _on_lowstate(msg) [background]
    GO2Connection->>GO2Connection: "_latest_lowstate = msg"

    Note over Caller,Robot: get_battery_soc()
    Caller->>GO2Connection: get_battery_soc()
    GO2Connection->>GO2Connection: _latest_lowstate["data"]["bms_state"]["soc"]
    GO2Connection-->>Caller: int (0-100) or None

    Note over Caller,Robot: set_rage_mode(enable)
    Caller->>GO2Connection: set_rage_mode(enable)
    GO2Connection->>WebRTC: set_rage_mode(enable)
    WebRTC->>Robot: balance_stand()
    WebRTC->>WebRTC: sleep(0.3s)
    WebRTC->>Robot: "2059 {data: enable}"
    alt "enable == True"
        WebRTC->>WebRTC: sleep(2.0s)
    end
    WebRTC->>Robot: SwitchJoystick(enable)
    WebRTC-->>GO2Connection: bool
    GO2Connection-->>Caller: bool

Loading

_{Reviews (11): Last reviewed commit: "Merge branch 'main' into ruthwik/fix/rag..." | Re-trigger Greptile}