diff --git a/fast_lowcmd_cpp/README.md b/fast_lowcmd_cpp/README.md new file mode 100644 index 0000000..15e3ba6 --- /dev/null +++ b/fast_lowcmd_cpp/README.md @@ -0,0 +1,45 @@ +# fast_lowcmd_cpp + +Native CPython extension for the Go1 PRO LowCmd hot path. + +It keeps the implementation separate from the main `go1_pro_sdk` package until +the robot-side timing and byte-equivalence checks are stable. + +## Build + +```bash +cd /root/go1_pro_sdk/fast_lowcmd_cpp +PYTHONPATH=/root/go1_pro_sdk python3 setup.py build_ext --inplace +``` + +On the Mac conda environment: + +```bash +cd /Users/chenyouyuan/cyy_ws/go1_pro_sdk/fast_lowcmd_cpp +PYTHONPATH=/Users/chenyouyuan/cyy_ws/go1_pro_sdk conda run -n free_dog_sdk python setup.py build_ext --inplace +``` + +## Test + +```bash +PYTHONPATH=.:.. python3 -m pytest test_fast_lowcmd.py +``` + +## Benchmark + +```bash +PYTHONPATH=.:.. python3 benchmark_fast_lowcmd.py +``` + +## API + +```python +from fast_lowcmd import FastLowCmdBuilder + +builder = FastLowCmdBuilder() +packet = builder.build_encrypted_servo12(q, kp=28.0, kd=0.7) +damping_packet = builder.build_encrypted_damping() +``` + +The generated bytes are expected to match the Python SDK exactly for the tested +packet forms. diff --git a/fast_lowcmd_cpp/benchmark_fast_lowcmd.py b/fast_lowcmd_cpp/benchmark_fast_lowcmd.py new file mode 100644 index 0000000..eaad975 --- /dev/null +++ b/fast_lowcmd_cpp/benchmark_fast_lowcmd.py @@ -0,0 +1,49 @@ +import statistics +import time + +from fast_lowcmd import FastLowCmdBuilder, default_state_path +from go1_pro_sdk import Blowfish, LowCmd, MotorCmd, MotorMode +from go1_pro_sdk.codec.lowcmd_builder import build_low_cmd_encrypted, build_low_cmd_plain + + +def bench(name, fn, n=5000): + for _ in range(100): + fn() + ts = [] + for _ in range(n): + t = time.perf_counter() + fn() + ts.append(time.perf_counter() - t) + ts.sort() + print( + f"{name:28s} median {statistics.median(ts) * 1e6:8.1f} us " + f"p95 {ts[int(.95 * (n - 1))] * 1e6:8.1f} us " + f"p99 {ts[int(.99 * (n - 1))] * 1e6:8.1f} us " + f"max {max(ts) * 1e6:8.1f} us" + ) + + +def main(): + q = [-0.1, 0.8, -1.5, 0.1, 0.8, -1.5, -0.1, 1.0, -1.5, 0.1, 1.0, -1.5] + + cmd = LowCmd() + for i in range(12): + cmd.set_motor(i, MotorCmd(mode=MotorMode.Servo, q=q[i], Kp=28.0, Kd=0.7)) + + bf = Blowfish.from_state_file(default_state_path()) + builder = FastLowCmdBuilder(default_state_path()) + + assert builder.build_plain_servo12(q, kp=28.0, kd=0.7) == build_low_cmd_plain(cmd) + assert builder.build_encrypted_servo12(q, kp=28.0, kd=0.7) == build_low_cmd_encrypted(cmd, bf) + + bench("python build_plain", lambda: build_low_cmd_plain(cmd), n=5000) + bench("python build_encrypted", lambda: build_low_cmd_encrypted(cmd, bf), n=2000) + bench("native build_plain_lowcmd", lambda: builder.build_plain_lowcmd(cmd), n=10000) + bench("native build_encrypted_lowcmd", lambda: builder.build_encrypted_lowcmd(cmd), n=10000) + bench("native build_plain_servo12", lambda: builder.build_plain_servo12(q, kp=28.0, kd=0.7), n=20000) + bench("native build_encrypted_servo12", lambda: builder.build_encrypted_servo12(q, kp=28.0, kd=0.7), n=10000) + bench("native damping encrypted", lambda: builder.build_encrypted_damping(), n=10000) + + +if __name__ == "__main__": + main() diff --git a/fast_lowcmd_cpp/fast_lowcmd.py b/fast_lowcmd_cpp/fast_lowcmd.py new file mode 100644 index 0000000..09974c1 --- /dev/null +++ b/fast_lowcmd_cpp/fast_lowcmd.py @@ -0,0 +1,49 @@ +"""Python wrapper for the native LowCmd builder. + +This module is intentionally kept outside the main go1_pro_sdk package until +the byte-equivalence and timing tests are stable on the robot board. +""" +from pathlib import Path + +from go1_fast_lowcmd import FastLowCmdBuilder as _NativeFastLowCmdBuilder + + +def default_state_path() -> str: + return str( + Path(__file__).resolve().parents[1] + / "go1_pro_sdk" + / "_data" + / "blowfish_state.bin" + ) + + +class FastLowCmdBuilder: + def __init__(self, state_path=None): + self._native = _NativeFastLowCmdBuilder(str(state_path or default_state_path())) + + def build_plain_damping(self): + return self._native.build_plain_damping() + + def build_encrypted_damping(self): + return self._native.build_encrypted_damping() + + def build_plain_servo12(self, q, dq=None, tau=None, kp=None, kd=None): + return self._native.build_plain_servo12(q, dq=dq, tau=tau, kp=kp, kd=kd) + + def build_encrypted_servo12(self, q, dq=None, tau=None, kp=None, kd=None): + return self._native.build_encrypted_servo12(q, dq=dq, tau=tau, kp=kp, kd=kd) + + def build_plain_fields20(self, mode=None, q=None, dq=None, tau=None, kp=None, kd=None): + return self._native.build_plain_fields20(mode=mode, q=q, dq=dq, tau=tau, kp=kp, kd=kd) + + def build_encrypted_fields20(self, mode=None, q=None, dq=None, tau=None, kp=None, kd=None): + return self._native.build_encrypted_fields20(mode=mode, q=q, dq=dq, tau=tau, kp=kp, kd=kd) + + def build_plain_lowcmd(self, lowcmd): + return self._native.build_plain_lowcmd(lowcmd) + + def build_encrypted_lowcmd(self, lowcmd): + return self._native.build_encrypted_lowcmd(lowcmd) + + def encrypt(self, data): + return self._native.encrypt(data) diff --git a/fast_lowcmd_cpp/go1_fast_lowcmd.cpp b/fast_lowcmd_cpp/go1_fast_lowcmd.cpp new file mode 100644 index 0000000..54f0299 --- /dev/null +++ b/fast_lowcmd_cpp/go1_fast_lowcmd.cpp @@ -0,0 +1,717 @@ +#define PY_SSIZE_T_CLEAN +#include +#include +#include +#include +#include +#include +#include + +namespace { + +constexpr Py_ssize_t LOWCMD_SIZE = 616; +constexpr Py_ssize_t CRC_OFFSET = 612; +constexpr Py_ssize_t MOTOR_OFFSET = 22; +constexpr Py_ssize_t MOTOR_SIZE = 27; +constexpr int MOTOR_COUNT = 20; +constexpr int SERVO_COUNT = 12; +constexpr uint32_t CRC_POLY = 0x04c11db7u; + +std::array make_crc_table() { + std::array table{}; + for (uint32_t byte = 0; byte < 256; ++byte) { + uint32_t crc = byte << 24; + for (int i = 0; i < 8; ++i) { + if (crc & 0x80000000u) { + crc = (crc << 1) ^ CRC_POLY; + } else { + crc <<= 1; + } + } + table[byte] = crc; + } + return table; +} + +const std::array CRC_TABLE = make_crc_table(); + +inline void put_u16_le(uint8_t* p, uint16_t v) { + p[0] = static_cast(v & 0xffu); + p[1] = static_cast((v >> 8) & 0xffu); +} + +inline void put_u32_le(uint8_t* p, uint32_t v) { + p[0] = static_cast(v & 0xffu); + p[1] = static_cast((v >> 8) & 0xffu); + p[2] = static_cast((v >> 16) & 0xffu); + p[3] = static_cast((v >> 24) & 0xffu); +} + +inline uint32_t get_u32_le(const uint8_t* p) { + return (static_cast(p[0])) + | (static_cast(p[1]) << 8) + | (static_cast(p[2]) << 16) + | (static_cast(p[3]) << 24); +} + +inline void put_f32_le(uint8_t* p, double value) { + float f = static_cast(value); + uint32_t raw; + static_assert(sizeof(raw) == sizeof(f), "float must be 32-bit"); + std::memcpy(&raw, &f, sizeof(raw)); + put_u32_le(p, raw); +} + +uint32_t crc32_unitree(const uint8_t* data, Py_ssize_t len) { + uint32_t crc = 0xffffffffu; + Py_ssize_t words = len / 4; + for (Py_ssize_t i = 0; i < words; ++i) { + uint32_t word = get_u32_le(data + i * 4); + crc = (crc << 8) ^ CRC_TABLE[((crc >> 24) ^ ((word >> 24) & 0xffu)) & 0xffu]; + crc = (crc << 8) ^ CRC_TABLE[((crc >> 24) ^ ((word >> 16) & 0xffu)) & 0xffu]; + crc = (crc << 8) ^ CRC_TABLE[((crc >> 24) ^ ((word >> 8) & 0xffu)) & 0xffu]; + crc = (crc << 8) ^ CRC_TABLE[((crc >> 24) ^ (word & 0xffu)) & 0xffu]; + } + return crc; +} + +double round_half_even(double value, double scale) { + double scaled = value * scale; + double floor_v = std::floor(scaled); + double frac = scaled - floor_v; + double rounded; + constexpr double eps = 1e-12; + if (frac > 0.5 + eps) { + rounded = floor_v + 1.0; + } else if (frac < 0.5 - eps) { + rounded = floor_v; + } else { + rounded = (std::fmod(floor_v, 2.0) == 0.0) ? floor_v : floor_v + 1.0; + } + return rounded / scale; +} + +void encode_tau(uint8_t* p, double tau) { + tau = round_half_even(tau, 100.0); + int integer_part = static_cast(tau); + double fractional_part = tau - static_cast(integer_part); + bool neg = false; + if (tau < 0.0) { + neg = true; + integer_part = 255 + integer_part; + } + if (fractional_part == 0.0) { + neg = false; + } + int hex_value = static_cast(fractional_part * 256.0); + if (neg) { + hex_value = 255 + hex_value + 1; + } + p[0] = static_cast(hex_value & 0xff); + p[1] = static_cast(integer_part & 0xff); +} + +void encode_kp(uint8_t* p, double kp) { + double base_d = std::floor(kp); + int base = static_cast(base_d); + double frac_d = round_half_even(kp - base_d, 10.0); + int frac = static_cast(frac_d * 10.0); + int val; + if (frac < 5) { + val = base * 32 + frac * 3; + } else { + val = base * 32 + (frac - 1) * 3 + 4; + } + put_u16_le(p, static_cast(val)); +} + +void encode_kd(uint8_t* p, double kd) { + int integer_part = static_cast(kd); + double fractional_part = round_half_even(kd - integer_part, 10.0); + int frac = static_cast(fractional_part * 10.0); + uint8_t nibble = 0; + switch (frac) { + case 0: nibble = 0x0; break; + case 1: nibble = 0x1; break; + case 2: nibble = 0x3; break; + case 3: nibble = 0x4; break; + case 4: nibble = 0x6; break; + case 5: nibble = 0x8; break; + case 6: nibble = 0x9; break; + case 7: nibble = 0xb; break; + case 8: nibble = 0xc; break; + case 9: nibble = 0xe; break; + default: nibble = 0x0; break; + } + put_u16_le(p, static_cast((integer_part << 4) | nibble)); +} + +void encode_motor(uint8_t* p, int mode, double q, double dq, double tau, double kp, double kd) { + p[0] = static_cast(mode & 0xff); + put_f32_le(p + 1, q); + put_f32_le(p + 5, dq); + encode_tau(p + 9, tau); + encode_kp(p + 11, kp); + encode_kd(p + 13, kd); + // reserve[0..2] are already zeroed by the packet initializer. +} + +void encode_motor_with_reserve( + uint8_t* p, int mode, double q, double dq, double tau, double kp, double kd, + uint32_t reserve0, uint32_t reserve1, uint32_t reserve2) { + encode_motor(p, mode, q, dq, tau, kp, kd); + put_u32_le(p + 15, reserve0); + put_u32_le(p + 19, reserve1); + put_u32_le(p + 23, reserve2); +} + +void init_lowcmd_header(uint8_t* cmd) { + std::memset(cmd, 0, LOWCMD_SIZE); + cmd[0] = 0xfe; + cmd[1] = 0xef; + cmd[2] = 0xff; + cmd[20] = 0x3a; + cmd[21] = 0xc0; +} + +void finish_crc(uint8_t* cmd) { + uint32_t crc = crc32_unitree(cmd, CRC_OFFSET); + put_u32_le(cmd + CRC_OFFSET, crc); +} + +bool read_double_sequence(PyObject* obj, double* out, int count, double default_value, const char* name) { + for (int i = 0; i < count; ++i) { + out[i] = default_value; + } + if (obj == nullptr || obj == Py_None) { + return true; + } + if (PyNumber_Check(obj) && !PySequence_Check(obj)) { + double v = PyFloat_AsDouble(obj); + if (PyErr_Occurred()) { + return false; + } + for (int i = 0; i < count; ++i) { + out[i] = v; + } + return true; + } + PyObject* seq = PySequence_Fast(obj, name); + if (!seq) { + return false; + } + Py_ssize_t n = PySequence_Fast_GET_SIZE(seq); + if (n < count) { + Py_DECREF(seq); + PyErr_Format(PyExc_ValueError, "%s must contain at least %d values", name, count); + return false; + } + PyObject** items = PySequence_Fast_ITEMS(seq); + for (int i = 0; i < count; ++i) { + out[i] = PyFloat_AsDouble(items[i]); + if (PyErr_Occurred()) { + Py_DECREF(seq); + return false; + } + } + Py_DECREF(seq); + return true; +} + +bool read_int_sequence(PyObject* obj, int* out, int count, int default_value, const char* name) { + for (int i = 0; i < count; ++i) { + out[i] = default_value; + } + if (obj == nullptr || obj == Py_None) { + return true; + } + if (PyLong_Check(obj)) { + long v = PyLong_AsLong(obj); + if (PyErr_Occurred()) { + return false; + } + for (int i = 0; i < count; ++i) { + out[i] = static_cast(v); + } + return true; + } + PyObject* seq = PySequence_Fast(obj, name); + if (!seq) { + return false; + } + Py_ssize_t n = PySequence_Fast_GET_SIZE(seq); + if (n < count) { + Py_DECREF(seq); + PyErr_Format(PyExc_ValueError, "%s must contain at least %d values", name, count); + return false; + } + PyObject** items = PySequence_Fast_ITEMS(seq); + for (int i = 0; i < count; ++i) { + long v = PyLong_AsLong(items[i]); + if (PyErr_Occurred()) { + Py_DECREF(seq); + return false; + } + out[i] = static_cast(v); + } + Py_DECREF(seq); + return true; +} + +bool get_int_attr(PyObject* obj, const char* name, int* out) { + PyObject* value = PyObject_GetAttrString(obj, name); + if (!value) { + return false; + } + long v = PyLong_AsLong(value); + Py_DECREF(value); + if (PyErr_Occurred()) { + return false; + } + *out = static_cast(v); + return true; +} + +bool get_double_attr(PyObject* obj, const char* name, double* out) { + PyObject* value = PyObject_GetAttrString(obj, name); + if (!value) { + return false; + } + double v = PyFloat_AsDouble(value); + Py_DECREF(value); + if (PyErr_Occurred()) { + return false; + } + *out = v; + return true; +} + +bool copy_bytes_attr(PyObject* obj, const char* name, uint8_t* out, Py_ssize_t expected_len) { + PyObject* value = PyObject_GetAttrString(obj, name); + if (!value) { + return false; + } + char* data = nullptr; + Py_ssize_t len = 0; + if (PyBytes_AsStringAndSize(value, &data, &len) < 0) { + Py_DECREF(value); + return false; + } + if (len != expected_len) { + Py_DECREF(value); + PyErr_Format(PyExc_ValueError, "%s must be %zd bytes", name, expected_len); + return false; + } + std::memcpy(out, data, static_cast(expected_len)); + Py_DECREF(value); + return true; +} + +bool read_reserve3(PyObject* motor, uint32_t reserve[3]) { + reserve[0] = reserve[1] = reserve[2] = 0; + PyObject* value = PyObject_GetAttrString(motor, "reserve"); + if (!value) { + return false; + } + PyObject* seq = PySequence_Fast(value, "motor.reserve must be a sequence"); + Py_DECREF(value); + if (!seq) { + return false; + } + Py_ssize_t n = PySequence_Fast_GET_SIZE(seq); + if (n < 3) { + Py_DECREF(seq); + PyErr_SetString(PyExc_ValueError, "motor.reserve must contain at least 3 values"); + return false; + } + PyObject** items = PySequence_Fast_ITEMS(seq); + for (int i = 0; i < 3; ++i) { + unsigned long v = PyLong_AsUnsignedLong(items[i]); + if (PyErr_Occurred()) { + Py_DECREF(seq); + return false; + } + reserve[i] = static_cast(v); + } + Py_DECREF(seq); + return true; +} + +PyObject* build_plain_from_lowcmd_obj(PyObject* lowcmd) { + PyObject* result = PyBytes_FromStringAndSize(nullptr, LOWCMD_SIZE); + if (!result) { + return nullptr; + } + auto* cmd = reinterpret_cast(PyBytes_AS_STRING(result)); + init_lowcmd_header(cmd); + + int value = 0; + if (!copy_bytes_attr(lowcmd, "head", cmd, 2)) { + Py_DECREF(result); + return nullptr; + } + if (!get_int_attr(lowcmd, "levelFlag", &value)) { + Py_DECREF(result); + return nullptr; + } + cmd[2] = static_cast(value & 0xff); + if (!get_int_attr(lowcmd, "frameReserve", &value)) { + Py_DECREF(result); + return nullptr; + } + cmd[3] = static_cast(value & 0xff); + if (!copy_bytes_attr(lowcmd, "SN", cmd + 4, 8)) { + Py_DECREF(result); + return nullptr; + } + if (!copy_bytes_attr(lowcmd, "version", cmd + 12, 8)) { + Py_DECREF(result); + return nullptr; + } + if (!get_int_attr(lowcmd, "bandWidth", &value)) { + Py_DECREF(result); + return nullptr; + } + cmd[20] = static_cast((value >> 8) & 0xff); + cmd[21] = static_cast(value & 0xff); + + PyObject* motors = PyObject_GetAttrString(lowcmd, "motorCmd"); + if (!motors) { + Py_DECREF(result); + return nullptr; + } + PyObject* seq = PySequence_Fast(motors, "lowcmd.motorCmd must be a sequence"); + Py_DECREF(motors); + if (!seq) { + Py_DECREF(result); + return nullptr; + } + if (PySequence_Fast_GET_SIZE(seq) < MOTOR_COUNT) { + Py_DECREF(seq); + Py_DECREF(result); + PyErr_SetString(PyExc_ValueError, "lowcmd.motorCmd must contain at least 20 motors"); + return nullptr; + } + PyObject** items = PySequence_Fast_ITEMS(seq); + for (int i = 0; i < MOTOR_COUNT; ++i) { + PyObject* motor = items[i]; + int mode = 0; + double q = 0.0, dq = 0.0, tau = 0.0, kp = 0.0, kd = 0.0; + uint32_t reserve[3]; + if (!get_int_attr(motor, "mode", &mode) + || !get_double_attr(motor, "q", &q) + || !get_double_attr(motor, "dq", &dq) + || !get_double_attr(motor, "tau", &tau) + || !get_double_attr(motor, "Kp", &kp) + || !get_double_attr(motor, "Kd", &kd) + || !read_reserve3(motor, reserve)) { + Py_DECREF(seq); + Py_DECREF(result); + return nullptr; + } + encode_motor_with_reserve( + cmd + MOTOR_OFFSET + i * MOTOR_SIZE, + mode, q, dq, tau, kp, kd, + reserve[0], reserve[1], reserve[2]); + } + Py_DECREF(seq); + + if (!copy_bytes_attr(lowcmd, "wirelessRemote", cmd + 566, 40)) { + Py_DECREF(result); + return nullptr; + } + if (!copy_bytes_attr(lowcmd, "reserve", cmd + 606, 4)) { + Py_DECREF(result); + return nullptr; + } + finish_crc(cmd); + return result; +} + +struct FastLowCmdBuilder { + PyObject_HEAD + uint32_t P[18]; + uint32_t S[4][256]; +}; + +uint32_t bf_f(FastLowCmdBuilder* self, uint32_t x) { + return (((self->S[0][(x >> 24) & 0xffu] + self->S[1][(x >> 16) & 0xffu]) & 0xffffffffu) + ^ self->S[2][(x >> 8) & 0xffu]) + self->S[3][x & 0xffu]; +} + +void encrypt_block_le(FastLowCmdBuilder* self, const uint8_t* in, uint8_t* out) { + uint32_t L = get_u32_le(in); + uint32_t R = get_u32_le(in + 4); + for (int i = 0; i < 16; ++i) { + L ^= self->P[i]; + R ^= bf_f(self, L); + uint32_t tmp = L; + L = R; + R = tmp; + } + uint32_t tmp = L; + L = R; + R = tmp; + R ^= self->P[16]; + L ^= self->P[17]; + put_u32_le(out, L); + put_u32_le(out + 4, R); +} + +PyObject* encrypt_bytes(FastLowCmdBuilder* self, const uint8_t* data, Py_ssize_t len) { + if (len % 8 != 0) { + PyErr_SetString(PyExc_ValueError, "data length must be a multiple of 8"); + return nullptr; + } + PyObject* result = PyBytes_FromStringAndSize(nullptr, len); + if (!result) { + return nullptr; + } + auto* out = reinterpret_cast(PyBytes_AS_STRING(result)); + for (Py_ssize_t i = 0; i < len; i += 8) { + encrypt_block_le(self, data + i, out + i); + } + return result; +} + +int FastLowCmdBuilder_init(FastLowCmdBuilder* self, PyObject* args, PyObject* kwargs) { + const char* state_path = nullptr; + static const char* kwlist[] = {"state_path", nullptr}; + if (!PyArg_ParseTupleAndKeywords(args, kwargs, "s", const_cast(kwlist), &state_path)) { + return -1; + } + + std::ifstream f(state_path, std::ios::binary); + if (!f) { + PyErr_Format(PyExc_FileNotFoundError, "cannot open Blowfish state file: %s", state_path); + return -1; + } + std::array buf{}; + f.read(reinterpret_cast(buf.data()), static_cast(buf.size())); + if (f.gcount() < static_cast(buf.size())) { + PyErr_Format(PyExc_ValueError, "Blowfish state file must contain at least 4168 bytes: %s", state_path); + return -1; + } + const uint8_t* p = buf.data(); + for (int i = 0; i < 18; ++i) { + self->P[i] = get_u32_le(p + i * 4); + } + p += 18 * 4; + for (int s = 0; s < 4; ++s) { + for (int i = 0; i < 256; ++i) { + self->S[s][i] = get_u32_le(p + (s * 256 + i) * 4); + } + } + return 0; +} + +PyObject* FastLowCmdBuilder_build_plain_damping(FastLowCmdBuilder*, PyObject*) { + PyObject* result = PyBytes_FromStringAndSize(nullptr, LOWCMD_SIZE); + if (!result) { + return nullptr; + } + auto* cmd = reinterpret_cast(PyBytes_AS_STRING(result)); + init_lowcmd_header(cmd); + finish_crc(cmd); + return result; +} + +PyObject* FastLowCmdBuilder_build_encrypted_damping(FastLowCmdBuilder* self, PyObject*) { + uint8_t cmd[LOWCMD_SIZE]; + init_lowcmd_header(cmd); + finish_crc(cmd); + return encrypt_bytes(self, cmd, LOWCMD_SIZE); +} + +PyObject* FastLowCmdBuilder_build_plain_servo12(FastLowCmdBuilder*, PyObject* args, PyObject* kwargs) { + PyObject* q_obj = nullptr; + PyObject* dq_obj = nullptr; + PyObject* tau_obj = nullptr; + PyObject* kp_obj = nullptr; + PyObject* kd_obj = nullptr; + static const char* kwlist[] = {"q", "dq", "tau", "kp", "kd", nullptr}; + if (!PyArg_ParseTupleAndKeywords( + args, kwargs, "O|OOOO", const_cast(kwlist), + &q_obj, &dq_obj, &tau_obj, &kp_obj, &kd_obj)) { + return nullptr; + } + + double q[SERVO_COUNT], dq[SERVO_COUNT], tau[SERVO_COUNT], kp[SERVO_COUNT], kd[SERVO_COUNT]; + if (!read_double_sequence(q_obj, q, SERVO_COUNT, 0.0, "q")) return nullptr; + if (!read_double_sequence(dq_obj, dq, SERVO_COUNT, 0.0, "dq")) return nullptr; + if (!read_double_sequence(tau_obj, tau, SERVO_COUNT, 0.0, "tau")) return nullptr; + if (!read_double_sequence(kp_obj, kp, SERVO_COUNT, 0.0, "kp")) return nullptr; + if (!read_double_sequence(kd_obj, kd, SERVO_COUNT, 0.0, "kd")) return nullptr; + + PyObject* result = PyBytes_FromStringAndSize(nullptr, LOWCMD_SIZE); + if (!result) { + return nullptr; + } + auto* cmd = reinterpret_cast(PyBytes_AS_STRING(result)); + init_lowcmd_header(cmd); + for (int i = 0; i < SERVO_COUNT; ++i) { + encode_motor(cmd + MOTOR_OFFSET + i * MOTOR_SIZE, 0x0a, q[i], dq[i], tau[i], kp[i], kd[i]); + } + finish_crc(cmd); + return result; +} + +PyObject* FastLowCmdBuilder_build_encrypted_servo12(FastLowCmdBuilder* self, PyObject* args, PyObject* kwargs) { + PyObject* plain = FastLowCmdBuilder_build_plain_servo12(self, args, kwargs); + if (!plain) { + return nullptr; + } + PyObject* result = encrypt_bytes( + self, + reinterpret_cast(PyBytes_AS_STRING(plain)), + PyBytes_GET_SIZE(plain)); + Py_DECREF(plain); + return result; +} + +PyObject* FastLowCmdBuilder_build_plain_fields20(FastLowCmdBuilder*, PyObject* args, PyObject* kwargs) { + PyObject* mode_obj = nullptr; + PyObject* q_obj = nullptr; + PyObject* dq_obj = nullptr; + PyObject* tau_obj = nullptr; + PyObject* kp_obj = nullptr; + PyObject* kd_obj = nullptr; + static const char* kwlist[] = {"mode", "q", "dq", "tau", "kp", "kd", nullptr}; + if (!PyArg_ParseTupleAndKeywords( + args, kwargs, "|OOOOOO", const_cast(kwlist), + &mode_obj, &q_obj, &dq_obj, &tau_obj, &kp_obj, &kd_obj)) { + return nullptr; + } + + int mode[MOTOR_COUNT]; + double q[MOTOR_COUNT], dq[MOTOR_COUNT], tau[MOTOR_COUNT], kp[MOTOR_COUNT], kd[MOTOR_COUNT]; + if (!read_int_sequence(mode_obj, mode, MOTOR_COUNT, 0, "mode")) return nullptr; + if (!read_double_sequence(q_obj, q, MOTOR_COUNT, 0.0, "q")) return nullptr; + if (!read_double_sequence(dq_obj, dq, MOTOR_COUNT, 0.0, "dq")) return nullptr; + if (!read_double_sequence(tau_obj, tau, MOTOR_COUNT, 0.0, "tau")) return nullptr; + if (!read_double_sequence(kp_obj, kp, MOTOR_COUNT, 0.0, "kp")) return nullptr; + if (!read_double_sequence(kd_obj, kd, MOTOR_COUNT, 0.0, "kd")) return nullptr; + + PyObject* result = PyBytes_FromStringAndSize(nullptr, LOWCMD_SIZE); + if (!result) { + return nullptr; + } + auto* cmd = reinterpret_cast(PyBytes_AS_STRING(result)); + init_lowcmd_header(cmd); + for (int i = 0; i < MOTOR_COUNT; ++i) { + encode_motor(cmd + MOTOR_OFFSET + i * MOTOR_SIZE, mode[i], q[i], dq[i], tau[i], kp[i], kd[i]); + } + finish_crc(cmd); + return result; +} + +PyObject* FastLowCmdBuilder_build_encrypted_fields20(FastLowCmdBuilder* self, PyObject* args, PyObject* kwargs) { + PyObject* plain = FastLowCmdBuilder_build_plain_fields20(self, args, kwargs); + if (!plain) { + return nullptr; + } + PyObject* result = encrypt_bytes( + self, + reinterpret_cast(PyBytes_AS_STRING(plain)), + PyBytes_GET_SIZE(plain)); + Py_DECREF(plain); + return result; +} + +PyObject* FastLowCmdBuilder_build_plain_lowcmd(FastLowCmdBuilder*, PyObject* args) { + PyObject* lowcmd = nullptr; + if (!PyArg_ParseTuple(args, "O", &lowcmd)) { + return nullptr; + } + return build_plain_from_lowcmd_obj(lowcmd); +} + +PyObject* FastLowCmdBuilder_build_encrypted_lowcmd(FastLowCmdBuilder* self, PyObject* args) { + PyObject* lowcmd = nullptr; + if (!PyArg_ParseTuple(args, "O", &lowcmd)) { + return nullptr; + } + PyObject* plain = build_plain_from_lowcmd_obj(lowcmd); + if (!plain) { + return nullptr; + } + PyObject* result = encrypt_bytes( + self, + reinterpret_cast(PyBytes_AS_STRING(plain)), + PyBytes_GET_SIZE(plain)); + Py_DECREF(plain); + return result; +} + +PyObject* FastLowCmdBuilder_encrypt(FastLowCmdBuilder* self, PyObject* args) { + Py_buffer view; + if (!PyArg_ParseTuple(args, "y*", &view)) { + return nullptr; + } + PyObject* result = encrypt_bytes(self, reinterpret_cast(view.buf), view.len); + PyBuffer_Release(&view); + return result; +} + +PyMethodDef FastLowCmdBuilder_methods[] = { + {"build_plain_damping", reinterpret_cast(FastLowCmdBuilder_build_plain_damping), METH_NOARGS, + "Build a plain all-damping LowCmd packet."}, + {"build_encrypted_damping", reinterpret_cast(FastLowCmdBuilder_build_encrypted_damping), METH_NOARGS, + "Build an encrypted all-damping LowCmd packet."}, + {"build_plain_servo12", reinterpret_cast(FastLowCmdBuilder_build_plain_servo12), + METH_VARARGS | METH_KEYWORDS, "Build a plain LowCmd packet with first 12 motors in servo mode."}, + {"build_encrypted_servo12", reinterpret_cast(FastLowCmdBuilder_build_encrypted_servo12), + METH_VARARGS | METH_KEYWORDS, "Build an encrypted LowCmd packet with first 12 motors in servo mode."}, + {"build_plain_fields20", reinterpret_cast(FastLowCmdBuilder_build_plain_fields20), + METH_VARARGS | METH_KEYWORDS, "Build a plain LowCmd packet from 20 motor field arrays."}, + {"build_encrypted_fields20", reinterpret_cast(FastLowCmdBuilder_build_encrypted_fields20), + METH_VARARGS | METH_KEYWORDS, "Build an encrypted LowCmd packet from 20 motor field arrays."}, + {"build_plain_lowcmd", reinterpret_cast(FastLowCmdBuilder_build_plain_lowcmd), METH_VARARGS, + "Build a plain LowCmd packet from a Python LowCmd object."}, + {"build_encrypted_lowcmd", reinterpret_cast(FastLowCmdBuilder_build_encrypted_lowcmd), METH_VARARGS, + "Build an encrypted LowCmd packet from a Python LowCmd object."}, + {"encrypt", reinterpret_cast(FastLowCmdBuilder_encrypt), METH_VARARGS, + "Encrypt a bytes-like object with the loaded Blowfish state."}, + {nullptr, nullptr, 0, nullptr}, +}; + +PyTypeObject FastLowCmdBuilderType = { + PyVarObject_HEAD_INIT(nullptr, 0) +}; + +PyModuleDef module = { + PyModuleDef_HEAD_INIT, + "go1_fast_lowcmd", + "Native CPython LowCmd builder/encrypter for go1_pro_sdk.", + -1, + nullptr, +}; + +} // namespace + +PyMODINIT_FUNC PyInit_go1_fast_lowcmd(void) { + FastLowCmdBuilderType.tp_name = "go1_fast_lowcmd.FastLowCmdBuilder"; + FastLowCmdBuilderType.tp_basicsize = sizeof(FastLowCmdBuilder); + FastLowCmdBuilderType.tp_flags = Py_TPFLAGS_DEFAULT; + FastLowCmdBuilderType.tp_doc = "Native LowCmd builder with persistent Blowfish state."; + FastLowCmdBuilderType.tp_init = reinterpret_cast(FastLowCmdBuilder_init); + FastLowCmdBuilderType.tp_new = PyType_GenericNew; + FastLowCmdBuilderType.tp_methods = FastLowCmdBuilder_methods; + + if (PyType_Ready(&FastLowCmdBuilderType) < 0) { + return nullptr; + } + + PyObject* m = PyModule_Create(&module); + if (!m) { + return nullptr; + } + Py_INCREF(&FastLowCmdBuilderType); + if (PyModule_AddObject(m, "FastLowCmdBuilder", reinterpret_cast(&FastLowCmdBuilderType)) < 0) { + Py_DECREF(&FastLowCmdBuilderType); + Py_DECREF(m); + return nullptr; + } + PyModule_AddIntConstant(m, "LOWCMD_SIZE", LOWCMD_SIZE); + return m; +} diff --git a/fast_lowcmd_cpp/setup.py b/fast_lowcmd_cpp/setup.py new file mode 100644 index 0000000..f612221 --- /dev/null +++ b/fast_lowcmd_cpp/setup.py @@ -0,0 +1,21 @@ +from pathlib import Path + +from setuptools import Extension, setup + + +ROOT = Path(__file__).resolve().parent + +setup( + name="go1_fast_lowcmd", + version="0.1.0", + description="Native CPython LowCmd builder/encrypter for go1_pro_sdk", + py_modules=["fast_lowcmd"], + ext_modules=[ + Extension( + "go1_fast_lowcmd", + sources=[str(ROOT / "go1_fast_lowcmd.cpp")], + language="c++", + extra_compile_args=["-std=c++17", "-O3"], + ) + ], +) diff --git a/fast_lowcmd_cpp/test_fast_lowcmd.py b/fast_lowcmd_cpp/test_fast_lowcmd.py new file mode 100644 index 0000000..c939b0d --- /dev/null +++ b/fast_lowcmd_cpp/test_fast_lowcmd.py @@ -0,0 +1,99 @@ +import os + +import pytest + +from fast_lowcmd import FastLowCmdBuilder, default_state_path +from go1_pro_sdk import Blowfish, LowCmd, MotorCmd, MotorMode +from go1_pro_sdk.codec.lowcmd_builder import build_low_cmd_encrypted, build_low_cmd_plain + + +@pytest.fixture +def builder(): + return FastLowCmdBuilder(default_state_path()) + + +@pytest.fixture +def bf(): + return Blowfish.from_state_file(default_state_path()) + + +def test_damping_plain_matches_python(builder): + assert builder.build_plain_damping() == build_low_cmd_plain(LowCmd()) + + +def test_damping_encrypted_matches_python(builder, bf): + cmd = LowCmd() + assert builder.build_encrypted_damping() == build_low_cmd_encrypted(cmd, bf) + + +def test_servo12_plain_matches_python(builder): + q = [0.1 * i - 0.5 for i in range(12)] + dq = [0.2 * i for i in range(12)] + tau = [0.0 for _ in range(12)] + kp = 36.0 + kd = 1.0 + + cmd = LowCmd() + for i in range(12): + cmd.set_motor(i, MotorCmd( + mode=MotorMode.Servo, + q=q[i], + dq=dq[i], + tau=tau[i], + Kp=kp, + Kd=kd, + )) + + assert builder.build_plain_servo12(q, dq=dq, tau=tau, kp=kp, kd=kd) == build_low_cmd_plain(cmd) + + +def test_servo12_encrypted_matches_python(builder, bf): + q = [-0.1, 0.8, -1.5, 0.1, 0.8, -1.5, -0.1, 1.0, -1.5, 0.1, 1.0, -1.5] + cmd = LowCmd() + for i in range(12): + cmd.set_motor(i, MotorCmd(mode=MotorMode.Servo, q=q[i], Kp=28.0, Kd=0.7)) + + assert builder.build_encrypted_servo12(q, kp=28.0, kd=0.7) == build_low_cmd_encrypted(cmd, bf) + + +def test_fields20_plain_matches_python(builder): + mode = [int(MotorMode.Servo) if i < 12 else int(MotorMode.Damping) for i in range(20)] + q = [0.05 * i for i in range(20)] + dq = [0.01 * i for i in range(20)] + tau = [0.0 for _ in range(20)] + kp = [20.0 if i < 12 else 0.0 for i in range(20)] + kd = [1.0 if i < 12 else 0.0 for i in range(20)] + + cmd = LowCmd() + for i in range(20): + cmd.set_motor(i, MotorCmd(mode=mode[i], q=q[i], dq=dq[i], tau=tau[i], Kp=kp[i], Kd=kd[i])) + + assert builder.build_plain_fields20(mode=mode, q=q, dq=dq, tau=tau, kp=kp, kd=kd) == build_low_cmd_plain(cmd) + + +def test_lowcmd_object_matches_python(builder, bf): + cmd = LowCmd() + cmd.frameReserve = 3 + cmd.wirelessRemote = bytes(range(40)) + cmd.reserve = b"\x01\x02\x03\x04" + for i in range(20): + cmd.set_motor( + i, + MotorCmd( + mode=MotorMode.Servo if i < 12 else MotorMode.Damping, + q=0.02 * i, + dq=0.01 * i, + tau=0.0, + Kp=12.0 if i < 12 else 0.0, + Kd=1.0 if i < 12 else 0.0, + reserve=[i, i + 1, i + 2], + ), + ) + + assert builder.build_plain_lowcmd(cmd) == build_low_cmd_plain(cmd) + assert builder.build_encrypted_lowcmd(cmd) == build_low_cmd_encrypted(cmd, bf) + + +def test_state_file_is_required(): + with pytest.raises(FileNotFoundError): + FastLowCmdBuilder(os.path.join(os.path.dirname(__file__), "missing_state.bin"))