feat(unitree): forward aes_128_key to WebRTC driver for G1 firmware >=1.5.1

[mihai-chiorean]

Motivation

Unitree G1 firmware 1.5.1 introduced a data2=3 WebRTC handshake that requires a per-device AES-128 key. Without it the connection fails inside unitree_webrtc_connect with RSA key format is not supported (pycryptodome trying to parse ciphertext bytes as an RSA key).

The per-device key is fetched from Unitree's cloud once via unitree-fetch-aes-key --email YOU --sn <serial> and then cached locally. The upstream driver legion1581/unitree_webrtc_connect@v2.1.1 accepts it through an aes_128_key= kwarg on UnitreeWebRTCConnection. The dimos wrapper at dimos/robot/unitree/connection.py never forwarded it, so anyone on G1 firmware >=1.5.1 cannot use the dimos Unitree integration today.

Change

UnitreeWebRTCConnection.__init__ now accepts an optional aes_128_key: str | None = None kwarg and falls back to the UNITREE_AES_128_KEY environment variable. The value is forwarded to the underlying LegionConnection only when non-empty, so the call is byte-identical to the previous behaviour when unset.

def __init__(self, ip: str, mode: str = "ai", aes_128_key: str | None = None) -> None:
    ...
    if aes_128_key is None:
        aes_128_key = os.environ.get("UNITREE_AES_128_KEY")
    extra: dict[str, Any] = {"aes_128_key": aes_128_key} if aes_128_key else {}
    self.conn = LegionConnection(WebRTCConnectionMethod.LocalSTA, ip=self.ip, **extra)

Verification

Tested on a Unitree G1 EDU+ (model string G1_Edu+) at firmware 1.5.1.1. With UNITREE_AES_128_KEY exported in the shell, dimos run unitree-g1-basic proceeds past the WebRTC handshake on 192.168.123.161:9991 and the robot accepts motion-switcher / wireless-controller commands as before. Without the env var, the same command reproduces RSA key format is not supported and the handshake never completes — same behaviour as main.

No breaking change

The kwarg is optional and defaults to None. The env-var lookup is opt-in (only consulted when the kwarg is None). Callers on G1 firmware <1.5.1 and all Go2 robots see no behavioural change — the **extra is empty and the LegionConnection call is byte-identical to the old one.

Optional follow-up (not in this PR)

pyproject.toml currently pins unitree-webrtc-connect-leshy>=2.0.7. The aes_128_key kwarg + the _decrypt_data1_v3 path that consumes it landed in legion1581/unitree_webrtc_connect@v2.1.1. If unitree-webrtc-connect-leshy is a downstream rebuild of that repo and is already on >=2.1.1, no change is needed and this PR is enough on its own. If it's still tracking 2.0.x, the dep will also need to be bumped (e.g. via [tool.uv.sources] pointing at git+https://github.com/legion1581/unitree_webrtc_connect.git@v2.1.1). Happy to follow up in a separate PR if maintainers confirm the situation.

[leshy]

reason for legion client fork is that unitree-webrtc-connect-leshy had a more efficinet lidar data parser and we had some issues with aiortc package pinning and had to fork it as well

legion1581/unitree_webrtc_connect@master...leshy:unitree_webrtc_connect:master
I assume lidar stuff is merged in the legion lib and all is resolved now

for FDEs (or anyone else here in the thread:) - we need to validate this on current g1s/go2s (ensure my lidar changes are now merged to legion client, transition to legion lib, test on our robots, make sure installation works, merge this)

[natcl]

We could also include GO2 firmware 1.1.15 as it has the same issue and this should fix it for the GO2 too !

[codecov]

Codecov Report

❌ Patch coverage is 33.33333% with 4 lines in your changes missing coverage. Please review.
✅ All tests successful. No failed tests found.

Files with missing lines	Patch %	Lines
dimos/robot/unitree/connection.py	33.33%	4 Missing ⚠️

📢 Thoughts on this report? Let us know!

[natcl]

had a more efficinet lidar data parser

I did notice a degradation in navigation performance on our go2 edu when I switched to the upstream library, might be worth checking if the changes are there.

[mihai-chiorean]

I Was at a hackathon, but I'll get this sorted out shortly.

[spomichter]

Thanks! @ruthwikdasyam or @Krishna can assist with any testing here

[mihai-chiorean]

Force-pushed 6bff60d44 addressing the open review items:

• @natcl — Go2 firmware ≥1.1.15 coverage added. aes_128_key now flows through go2/connection.py (ConnectionConfig field → make_connection(..., aes_128_key=...) → UnitreeWebRTCConnection) mirroring the G1 pattern, plus through go2/fleet_connection.py so fleet followers get the same forwarding (with an inline comment noting per-device-key heterogeneous fleets are a future per-IP-mapping follow-up).
• @greptile-apps — The previous if aes_128_key is None → if not aes_128_key truthiness fix landed in fceccd7c0. Updated the now-misleading docstring on test_empty_string_kwarg_* and added a sibling test_empty_string_kwarg_uses_env_when_set covering the YAML/JSON aes_128_key: "" regression you flagged.
• Added a targeted unit test go2/test_connection.py::test_make_connection_webrtc_forwards_aes_128_key to pin the new make_connection ↔ UnitreeWebRTCConnection contract.

Ran the changes through codex review --commit HEAD after each iteration. Codex caught the missed Go2FleetConnection.start() forwarding (now fixed) and signed off on the amended commit.

@leshy — re. the upstream legion lib transition + lidar perf, leaving that for a separate PR (also flagged on this PR thread by @natcl who noticed degraded nav on Go2 EDU after the upstream switch). Happy to scope that follow-up once you confirm the lidar parser changes have landed in legion upstream.

[leshy]

I think ideally should bump to latest unitree_webrtc_connect for this to be usable.

my lidar changes have landed upstream looks like (legion1581/unitree_webrtc_connect#46) but let's still compare perforamnce between https://github.com/leshy/unitree_webrtc_connect (what we use now) and upgraded latest official unitree_webrtc_connect

because of

I did notice a degradation in navigation performance on our go2 edu when I switched to the upstream library, might be worth checking if the changes are there.

@ruthwikdasyam has experience with lidar data profiling so can look into this. if performance is the same, we can commit a bump in pyproject to official lib and merge all together

[mihai-chiorean]

I think ideally should bump to latest unitree_webrtc_connect for this to be usable.

my lidar changes have landed upstream looks like (legion1581/unitree_webrtc_connect#46) but let's still compare perforamnce between https://github.com/leshy/unitree_webrtc_connect (what we use now) and upgraded latest official unitree_webrtc_connect

because of

I did notice a degradation in navigation performance on our go2 edu when I switched to the upstream library, might be worth checking if the changes are there.

@ruthwikdasyam has experience with lidar data profiling so can look into this. if performance is the same, we can commit a bump in pyproject to official lib and merge all together

want me to do that here or is that for a separate PR...? happy to update it, a little worried about the blast radius of updating a package like that since I'm new to the repo.

[greptile-apps]

Greptile Summary

This PR adds optional aes_128_key support to the Unitree WebRTC connection layer, required by G1 firmware ≥1.5.1 and Go2 firmware ≥1.1.15 for a new data2=3 handshake. The key can be supplied via a new constructor kwarg or via the UNITREE_AES_128_KEY environment variable, and falls back gracefully to the previous code path when unset.

• UnitreeWebRTCConnection.__init__ gains an optional aes_128_key kwarg; a not aes_128_key truthiness guard (rather than is None) ensures both None and "" fall through to the env-var lookup, and the key is only injected into LegionConnection when truthy.
• G1Config, G1HighLevelWebRtcConfig, and ConnectionConfig (Go2) each gain an aes_128_key: str | None = None field wired through to the connection constructor; Go2FleetConnection.start forwards the same configured key to all follower robots (a per-IP mapping is noted as future work).
• Comprehensive unit tests cover the no-key default, explicit-kwarg, env-var-fallback, kwarg-over-env precedence, and empty-string-falls-through-to-env-var scenarios.

Confidence Score: 5/5

Safe to merge — all call sites are updated, the new kwarg is optional with a backward-compatible default, and the truthiness guard correctly handles both None and "" inputs.

The change is narrow and well-contained: a single new optional parameter threads through four config classes and the connection factory, each wired consistently. The previously reported empty-string bypass was resolved before this review, and the new test suite explicitly covers that scenario along with the no-key, explicit-kwarg, env-var-fallback, and kwarg-over-env-precedence cases. There are no logic gaps or missing call sites.

No files require special attention.

Important Files Changed

Filename	Overview
dimos/robot/unitree/connection.py	Core change: adds aes_128_key kwarg with truthiness-guarded env-var fallback; correctly omits the key from LegionConnection when unset.
dimos/robot/unitree/g1/connection.py	Adds aes_128_key to G1Config and forwards it to UnitreeWebRTCConnection; straightforward wiring change.
dimos/robot/unitree/go2/connection.py	Adds aes_128_key to ConnectionConfig, threads it through make_connection, and into UnitreeWebRTCConnection; clean implementation.
dimos/robot/unitree/go2/fleet_connection.py	Fleet followers receive the same configured key as the primary; the single-key limitation for heterogeneous fleets is documented in a comment.
dimos/robot/unitree/test_connection.py	New unit tests cover all key-forwarding scenarios including the empty-string edge case; test design is correct.
dimos/robot/unitree/go2/test_connection.py	New test pins the Go2 make_connection → UnitreeWebRTCConnection kwarg wiring; straightforward and correct.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
    A[Caller provides config] --> B{aes_128_key in config?}
    B -- Yes --> C[Forward kwarg to UnitreeWebRTCConnection]
    B -- No --> C2[Pass None to UnitreeWebRTCConnection]
    C --> D{if not aes_128_key}
    C2 --> D
    D -- False: kwarg is truthy --> F[Skip env-var lookup]
    D -- True: kwarg is None or empty string --> E[Read UNITREE_AES_128_KEY from env]
    E --> G{env var set?}
    G -- Yes --> H[Use env var value]
    G -- No --> I[key remains None]
    F --> J{Final key truthy?}
    H --> J
    I --> J
    J -- Yes --> K[LegionConnection called with key kwarg]
    J -- No --> L[LegionConnection called without key kwarg - pre-PR behavior]

Loading

_{Reviews (8): Last reviewed commit: "Merge branch 'main' into feat/forward-ae..." | Re-trigger Greptile}