Trossen Arm MuJoCo自定义2：随机化物体起始位置

实现方法很简单在任务类的initialize_episode方法中为物体的自由关节free joint生成随机的位姿位置 朝向并赋值给仿真中的qpos下面以单臂抓取葡萄串为例展示如何让葡萄串在每个 episode 开始时出现在工作空间内的随机位置1. 修改ee_sim_env.py中的任务类找到你正在使用的任务类例如PickGrapeTask修改其initialize_episode方法import numpy as np from scipy.spatial.transform import Rotation as R class PickGrapeTask(TrossenAIWXAIEETask): def initialize_episode(self, physics: Physics) - None: # 1. 重置机械臂到初始姿态继承方法 self.initialize_robots(physics) # 2. 生成随机位置在夹爪可达范围内 x np.random.uniform(0.2, 0.45) # 前后方向米 y np.random.uniform(-0.25, 0.25) # 左右方向 z 0.05 # 桌面高度固定假设物体底面接触桌面 # 3. 生成随机朝向绕竖直轴旋转也可以完全随机 yaw np.random.uniform(-np.pi, np.pi) quat R.from_euler(z, yaw).as_quat() # 返回 [x, y, z, w] 格式 # 组合成7维位姿 [x, y, z, qw, qx, qy, qz] (注意文档中顺序为 [x,y,z,w,x,y,z]) # 这里使用 MuJoCo 标准顺序位置 四元数 [w, x, y, z] random_pose np.array([x, y, z, quat[3], quat[0], quat[1], quat[2]]) # 4. 获取物体关节的索引 grape_joint_idx physics.model.name2id(grape_bunch_joint, joint) # 5. 在重置上下文中赋值避免与仿真步进冲突 with physics.reset_context(): np.copyto(physics.data.qpos[grape_joint_idx:grape_joint_idx7], random_pose) # 6. 调用父类初始化如有必要 super().initialize_episode(physics)位置范围需要根据机械臂的工作空间和相机视野调整。上例中x0.2~0.45y-0.25~0.25适用于 WidowX AI 在桌面上的典型工作区域。z 坐标假设物体放置在桌面上桌面高度在 XML 中定义为0.05米如果物体的碰撞几何体中心在物体几何中心则z应设为桌面高度 物体半高。你也可以通过读取桌面高度并加上物体尺寸动态计算。四元数顺序MuJoCo 的qpos中四元数顺序是[w, x, y, z]而scipy.spatial.transform.Rotation返回的是[x, y, z, w]所以赋值时需要调换顺序。随机种子如果需要实验可重复可以在脚本开头设置np.random.seed(0)。2. 保持get_env_state返回随机位姿任务类中的get_env_state方法应该返回物体的当前位姿通常就是physics.data.qpos中物体关节的那一段。这不需要额外修改因为已经自动反映了随机化后的值。staticmethod def get_env_state(physics: Physics) - np.ndarray: # 返回葡萄串的7维位姿记录到数据集中 grape_joint_idx physics.model.name2id(grape_bunch_joint, joint) return physics.data.qpos[grape_joint_idx:grape_joint_idx7].copy()3. 同步修改sim_env.py中的任务类重放阶段在sim_env.py中对应的任务类例如PickGrapeTaskinitialize_episode方法不需要再随机化位置而是应该使用从数据集传入的固定位置由self.grape_pose等变量接收。框架会自动将ee_sim_env中记录的环境状态传递给sim_env的任务实例。只需要确保sim_env中的任务类有一个属性来接收该位姿并在initialize_episode中应用它class PickGrapeTask(TrossenAIWXAITask): def __init__(self, physics, **kwargs): super().__init__(physics, **kwargs) self.grape_pose None # 将由 record 脚本自动填充 def initialize_episode(self, physics: Physics) - None: with physics.reset_context(): # 重置机械臂 physics.named.data.qpos[:7] START_ARM_POSE # 设置葡萄串位置如果存在传入的位姿 if self.grape_pose is not None: grape_joint_idx physics.model.name2id(grape_bunch_joint, joint) np.copyto(physics.data.qpos[grape_joint_idx:grape_joint_idx7], self.grape_pose) super().initialize_episode(physics)4. 随机化多个物体或额外参数如果你有多个物体需要随机化比如两个方块、一个障碍物可以在initialize_episode中为每个物体生成独立的随机位姿并分别赋值到对应的关节索引。对于其他参数如灯光颜色、相机噪声也可以类似地随机化但需要确保在get_env_state中记录这些参数可以扩展返回的数组并在sim_env中恢复5. 验证随机化效果python trossen_arm_mujoco/scripts/record_sim_episodes.py \ --task_name pick_grape \ --data_dir test_random \ --num_episodes 1 \ --onscreen_render你会看到每次运行时葡萄串出现在不同的位置。如果希望在不运行完整数据收集的情况下快速测试可以直接运行ee_sim_env.py中的任务测试代码如有。