SUMO:动态全身运动操作系统的技术深度解析
1. 研究背景与核心问题
在机器人领域,让腿足机器人完成动态物体操作一直是一个极具挑战性的问题。传统的机器人系统往往将运动控制和物体操作分离处理,这导致机器人在面对大型、重型物体时表现不佳。SUMO(Dynamic and Generalizable Whole-Body Loco-Manipulation)项目由MIT、RAI Institute和Cornell联合研发,提出了一种创新的全身协同控制方法,使得Spot四足机器人能够操作超过其额定负载能力的物体,例如将重达15公斤的轮胎竖立起来(而Spot的峰值提升能力仅为11公斤),或者拖动比自身更大的人群控制栅栏。这项研究的核心洞察在于:通过在测试时使用基于采样的规划器对预训练的全身控制策略进行引导,可以使机器人解决各种动态运动操作任务,并且无需额外调整或训练即可泛化到不同的物体和任务。
2. 系统架构设计
SUMO采用分层控制架构,将复杂的全身运动操作问题分解为两个层次:高层采样式模型预测控制(MPC)和低层全身控制策略。这种设计充分利用了各层的优势,高层MPC负责任务级规划,以20Hz的频率生成期望的躯干、手臂和腿部指令;低层策略则负责将这些高级指令转换为关节级控制命令,以50Hz的频率执行。这种分层结构不仅简化了控制问题的复杂度,还使得系统具有良好的泛化能力。与端到端强化学习方法相比,SUMO在面对不同物体时表现出更强的鲁棒性,因为它可以在测试时通过替换物体模型或调整代价函数来适应新任务,而无需重新训练神经网络。

下图展示了SUMO的分层控制流程,清晰地说明了高层规划和低层控制之间的交互关系:
2.1 分层控制的数学表述
在SUMO的框架中,控制问题被形式化为一个分层优化过程。高层MPC求解器在每个时间步 t t t优化以下目标函数:
min a t : t + H ∑ k = t t + H c ( s k , a k ) \min_{a_{t:t+H}} \sum_{k=t}^{t+H} c(s_k, a_k) at:t+Hmink=t∑t+Hc(sk,ak)
其中 a k a_k ak表示高级动作(躯干速度、手臂位置等), s k s_k sk是状态, c ( ⋅ ) c(\cdot) c(⋅)是任务相关的代价函数, H H H是预测时域。关键的创新点在于,状态转移不是通过标准的多体动力学模型 s k + 1 = f ( s k , u k ) s_{k+1} = f(s_k, u_k) sk+1=f(sk,uk)计算(其中 u k u_k uk是关节力矩),而是通过神经网络策略增强的动力学模型 s k + 1 = f ( s k , π ( s k , a k ) ) s_{k+1} = f(s_k, \pi(s_k, a_k)) sk+1=f(sk,π(sk,ak))实现,这里 π \pi π是预训练的低层策略。这种设计使得MPC可以在更抽象的动作空间中进行规划,大大降低了优化问题的维度和复杂度。
3. 代码实现架构
SUMO项目基于Judo框架构建,采用模块化设计,主要包含以下核心组件:
3.1 项目结构
sumo/
├── sumo/ # 主Python包
│ ├── tasks/ # 任务实现
│ │ ├── g1/ # G1人形机器人任务
│ │ └── spot/ # Spot四足机器人任务
│ ├── controller/ # 控制器实现
│ ├── models/ # MuJoCo XML和网格资源
│ ├── cli.py # 命令行入口
│ └── run_mpc/ # 无头MPC运行器
├── g1_extensions/ # G1原生扩展(C++/pybind11)
└── tests/ # 测试套件
下图展示了完整的项目结构,包括Python包、C++扩展和配置文件的组织方式:
3.2 任务配置系统
SUMO使用数据类(dataclass)定义任务配置,这种设计使得参数管理清晰且类型安全。以Spot推箱子任务为例:
@dataclass
class SpotBoxPushConfig(SpotPushConfig):
"""Spot推箱子任务的配置"""
goal_position: np.ndarray = np_1d_field(
np.array([0.0, 0.0, BOX_HALF_LENGTH]),
names=["x", "y", "z"],
mins=[-5.0, -5.0, 0.0],
maxs=[5.0, 5.0, 3.0],
vis_name="goal_position",
xyz_vis_indices=[0, 1, None],
)
这个配置类继承自SpotPushConfig,定义了目标位置参数。np_1d_field是一个自定义字段工具,不仅存储数值,还包含了参数的语义信息(名称、范围、可视化索引),这对于交互式调试和参数调优非常有用。
3.3 奖励函数设计
奖励函数是强化学习和MPC系统的核心。SUMO采用模块化的奖励函数设计,将复杂的任务目标分解为多个可组合的奖励项。以Spot推箱子任务为例,其奖励函数实现如下:
def reward(
self,
states: np.ndarray,
sensors: np.ndarray,
controls: np.ndarray,
system_metadata: dict[str, Any] | None = None,
) -> np.ndarray:
"""使用目标距离、夹爪距离和物体速度计算奖励"""
batch_size = states.shape[0]
qpos = states[..., : self.model.nq]
object_pos = qpos[..., self.object_pose_idx : self.object_pose_idx + 3]
gripper_pos = sensors[..., self.gripper_pos_idx : self.gripper_pos_idx + 3]
object_linear_velocity = states[..., self.object_vel_idx : self.object_vel_idx + 3]
goal_reward = goal_distance_reward(self.config, object_pos)
gripper_proximity_reward = gripper_distance_reward(
self.config, np.linalg.norm(gripper_pos - object_pos, axis=-1)
)
object_linear_velocity_penalty = object_linear_velocity_reward(
self.config, object_linear_velocity
)
return goal_reward + gripper_proximity_reward + object_linear_velocity_penalty
这个奖励函数包含三个关键组成部分:目标距离奖励鼓励物体接近目标位置,夹爪接近度奖励促使机器人保持与物体的接触,物体速度惩罚则防止物体运动过快导致失控。每个奖励项都有对应的权重参数(w_goal、w_gripper_proximity、w_object_velocity),可以在测试时灵活调整以适应不同的任务需求。这种设计体现了SUMO的核心优势:通过调整代价函数而非重新训练模型来实现任务泛化。
下图展示了奖励函数的设计架构,说明了从任务配置到最终优化目标的完整流程:
3.4 G1人形机器人任务实现
对于G1人形机器人,任务实现更加复杂,因为需要考虑双足平衡、手臂协调和全身姿态控制。以G1推箱子任务为例:
@dataclass
class G1BoxConfig(G1BaseConfig):
"""G1推箱子任务配置"""
goal_position: np.ndarray = np_1d_field(
np.array([0.0, 0.0, 0.0], dtype=np.float64),
names=["x", "y", "z"],
mins=[-5.0, -5.0, 0.0],
maxs=[5.0, 5.0, 1.0],
steps=[0.1, 0.1, 0.05],
vis_name="box_goal_position",
xyz_vis_indices=[0, 1, 2],
)
w_goal: float = 50.0
w_orientation: float = 15.0
w_pelvis_proximity: float = 0.1
w_hand_proximity: float = 10.0
w_robot_orientation: float = 50.0
w_controls: float = 3
G1的奖励函数需要平衡更多的约束条件,包括目标距离、箱子姿态、骨盆接近度、手部接近度、机器人朝向和控制代价。特别值得注意的是机器人朝向奖励的实现:
# 从四元数计算前向方向(X轴)
# 对于四元数[w, x, y, z],旋转后的前向向量(原始为[1, 0, 0])为:
# forward_x = 1 - 2(y^2 + z^2)
body_quat = qpos[..., self.body_pose_idx[3:7]]
y_q = body_quat[..., 2]
z_q = body_quat[..., 3]
forward_x = 1 - 2 * (y_q**2 + z_q**2)
# 当forward_x为正时给予奖励(面向+X方向)
robot_orientation_reward = config.w_robot_orientation * forward_x.mean(-1)
这段代码通过四元数直接计算机器人的前向方向,避免了昂贵的旋转矩阵转换。当机器人面向前方(默认朝向)时,forward_x接近1,获得最大奖励;当机器人背对目标时,forward_x为负,受到惩罚。这种设计确保机器人在推动物体时保持合理的身体朝向。
4. 高性能C++扩展实现
为了实现高效的策略推理和动力学仿真,SUMO为G1机器人开发了原生C++扩展,使用pybind11与Python接口。这个扩展的核心是基于ONNX Runtime的策略推理引擎。
4.1 ONNX策略包装器
class OnnxPolicy {
public:
explicit OnnxPolicy(const std::shared_ptr<Ort::Session>& session)
: session_(session),
memory_info_(Ort::MemoryInfo::CreateCpu(OrtDeviceAllocator, OrtMemTypeCPU)) {
Ort::Allocator allocator(*session_, memory_info_);
input_name_ = session_->GetInputNameAllocated(0, allocator).get();
output_name_ = session_->GetOutputNameAllocated(0, allocator).get();
input_shape_ = session_->GetInputTypeInfo(0).GetTensorTypeAndShapeInfo().GetShape();
output_shape_ = session_->GetOutputTypeInfo(0).GetTensorTypeAndShapeInfo().GetShape();
input_size_ = static_cast<int>(input_shape_[1]);
output_size_ = static_cast<int>(output_shape_[1]);
}
std::vector<float> run(const std::vector<float>& observation) {
if ((int)observation.size() != input_size_) {
throw std::runtime_error("观测维度与ONNX输入维度不匹配");
}
std::array<int64_t, 2> ishape = { 1, static_cast<int64_t>(observation.size()) };
Ort::Value input_tensor = Ort::Value::CreateTensor<float>(
memory_info_, const_cast<float*>(observation.data()),
observation.size(), ishape.data(), 2);
const char* in_names[1] = { input_name_.c_str() };
const char* out_names[1] = { output_name_.c_str() };
auto outputs = session_->Run(run_options_, in_names, &input_tensor, 1, out_names, 1);
auto& out = outputs[0];
float* ptr = out.GetTensorMutableData<float>();
auto info = out.GetTensorTypeAndShapeInfo();
size_t n = info.GetElementCount();
return std::vector<float>(ptr, ptr + n);
}
private:
std::shared_ptr<Ort::Session> session_;
Ort::MemoryInfo memory_info_;
Ort::RunOptions run_options_;
std::string input_name_;
std::string output_name_;
std::vector<int64_t> input_shape_;
std::vector<int64_t> output_shape_;
int input_size_ = 0;
int output_size_ = 0;
};
这个类封装了ONNX Runtime的会话管理和推理过程。在构造函数中,它自动提取模型的输入输出信息,包括张量名称和形状。run方法接收观测向量,创建输入张量,执行推理,并返回动作向量。使用ONNX格式的优势在于:它是一个开放标准,支持多种深度学习框架导出的模型;ONNX Runtime经过高度优化,在CPU上也能实现高效推理;模型文件独立于训练框架,便于部署和版本管理。
4.2 会话分配与优化配置
static std::shared_ptr<Ort::Session> allocate_shared_session(const std::string& onnx_path) {
static Ort::Env env;
Ort::SessionOptions opts;
opts.SetIntraOpNumThreads(1);
opts.SetGraphOptimizationLevel(GraphOptimizationLevel::ORT_ENABLE_BASIC);
return std::make_shared<Ort::Session>(env, onnx_path.c_str(), opts);
}
这个函数负责创建ONNX Runtime会话。关键配置包括:将线程内并行度设置为1(SetIntraOpNumThreads(1)),这在MPC的批量推理场景中可以避免线程竞争;启用基本图优化(ORT_ENABLE_BASIC),包括常量折叠、冗余节点消除等,在不显著增加加载时间的前提下提升推理性能。使用std::shared_ptr管理会话生命周期,确保多个策略实例可以安全共享同一个模型。
5. 控制器与优化器集成
SUMO的控制器模块负责协调高层MPC和低层策略的交互。控制器支持多种优化算法,包括交叉熵方法(CEM)和模型预测路径积分(MPPI)。
5.1 控制器工厂函数
def make_controller(
init_task: str,
init_optimizer: str,
task_registration_cfg: DictConfig | None = None,
optimizer_registration_cfg: DictConfig | None = None,
rollout_backend: Literal["mujoco"] = "mujoco",
) -> Controller:
"""创建支持G1后端的控制器"""
return _judo_make_controller(
init_task=init_task,
init_optimizer=init_optimizer,
task_registration_cfg=task_registration_cfg,
optimizer_registration_cfg=optimizer_registration_cfg,
rollout_backend=rollout_backend,
custom_rollout_backends={"mujoco_g1": G1RolloutBackend},
)
这个工厂函数是SUMO对Judo框架的扩展点。通过custom_rollout_backends参数注册自定义的G1推演后端,使得控制器可以使用C++加速的策略推理。这种设计遵循了开闭原则:对扩展开放(可以添加新的后端),对修改封闭(不需要改动Judo的核心代码)。
5.2 采样式MPC的工作原理
采样式MPC是SUMO高层规划的核心算法。与传统的基于梯度的优化方法不同,采样式MPC通过在动作空间中采样大量轨迹,评估每条轨迹的代价,然后根据代价对采样分布进行更新。以交叉熵方法(CEM)为例,其迭代过程如下:
- 从当前分布(通常是高斯分布)中采样N条动作序列
- 对每条序列进行前向推演,计算累积代价
- 选择代价最低的K条序列(精英样本)
- 用精英样本的均值和方差更新采样分布
- 重复步骤1-4直到收敛或达到最大迭代次数
这种方法的优势在于:不需要计算梯度,适用于非光滑的代价函数和动力学模型;可以自然地处理约束条件;易于并行化,适合GPU加速。在SUMO中,采样式MPC与神经网络策略的结合使得系统既能利用学习到的运动先验,又能在线适应新的任务目标。
下图详细展示了交叉熵方法的完整优化流程:
这个流程图清楚地展示了CEM算法的四个关键阶段:采样阶段生成候选解,推演阶段评估每个候选解的性能,评估阶段筛选出最优解,更新阶段改进采样分布。这个过程不断迭代,直到找到满意的解或达到计算预算限制。
6. 环境配置与构建系统
SUMO使用Pixi作为环境管理器,这是一个基于Conda的现代化包管理工具,支持跨平台依赖管理和任务自动化。
6.1 依赖管理
项目的核心依赖定义在pyproject.toml中:
[project]
name = "sumo"
version = "0.0.1"
dependencies = [
"judo-rai @ git+https://github.com/bdaiinstitute/judo.git@dta/fix_for_sumo",
"numpy",
"mujoco>=3.5.0, <3.6.0",
"h5py",
"tyro",
"tqdm",
"scipy",
]
[tool.pixi.feature.cpp.dependencies]
cmake = "*"
ninja = "*"
pybind11 = "*"
cxx-compiler = "*"
c-compiler = "*"
eigen = "*"
llvm-openmp = "*"
这里值得注意的是MuJoCo版本的严格限制(>=3.5.0, <3.6.0),这是因为MuJoCo的API在不同版本间可能存在不兼容变化。C++特性依赖包括CMake构建系统、Ninja构建工具、pybind11绑定库、Eigen线性代数库和OpenMP并行计算库。这些依赖通过Conda统一管理,确保在不同平台上的一致性。
6.2 构建流程
…详情请参照古月居
openEuler 是由开放原子开源基金会孵化的全场景开源操作系统项目,面向数字基础设施四大核心场景(服务器、云计算、边缘计算、嵌入式),全面支持 ARM、x86、RISC-V、loongArch、PowerPC、SW-64 等多样性计算架构
更多推荐

所有评论(0)