如何在PyBullet中模拟物体平衡（集成OpenAI Gym）

看起来你已经有了不错的基础框架，接下来我会帮你整合随机形状生成、物理控制和Gym环境逻辑，解决物体平衡的核心问题。

一、先修正现有Gym环境的基础Bug #

你的BeamEnv里有几个影响运行的小问题，先处理掉：

1. 拼写错误：gym.spaces.Descrete → gym.spaces.Discrete
2. 缺失导入：需要添加 from gym import spaces 和 import gym.utils.seeding as seeding
3. reset方法存在两段重复逻辑，需要合并PyBullet初始化与环境状态重置
4. _compute_observation中的未定义变量（如self.vt），需替换为实际物理观测值

二、实现随机形状生成 #

PyBullet支持两种随机形状生成方式，你可以根据需求选择：

方式1：动态生成随机凸多边形 #

适合快速生成简单随机形状，无需预存模型文件：

def create_random_convex_shape(mass=1.0, position=[0,0,1]):
    # 生成3-7个顶点的随机凸多边形
    num_vertices = np.random.randint(3, 8)
    angles = np.linspace(0, 2*np.pi, num_vertices, endpoint=False)
    radii = np.random.uniform(0.2, 0.5, num_vertices)
    vertices = np.array([[r*np.cos(a), r*np.sin(a), 0] for r,a in zip(radii, angles)])
    
    # 创建碰撞与视觉形状
    collision_shape = pb.createCollisionShape(
        shapeType=pb.GEOM_MESH,
        vertices=vertices.tolist()
    )
    visual_shape = pb.createVisualShape(
        shapeType=pb.GEOM_MESH,
        vertices=vertices.tolist(),
        rgbaColor=[np.random.rand(), np.random.rand(), np.random.rand(), 1]
    )
    
    # 创建物理体
    body_id = pb.createMultiBody(
        baseMass=mass,
        baseCollisionShapeIndex=collision_shape,
        baseVisualShapeIndex=visual_shape,
        basePosition=position
    )
    return body_id

方式2：随机加载预生成的.obj文件 #

如果你有一批模型文件存在procedural_objects目录下，可以随机选择加载：

import glob

def load_random_obj(mass=1.0, position=[0,0,1], scale=[0.1,0.1,0.1]):
    obj_files = glob.glob("procedural_objects/**/*.obj", recursive=True)
    if not obj_files:
        #  fallback到随机凸多边形
        return create_random_convex_shape(mass, position)
    
    random_obj = np.random.choice(obj_files)
    collision_shape = pb.createCollisionShape(
        shapeType=pb.GEOM_MESH,
        fileName=random_obj,
        meshScale=scale
    )
    visual_shape = pb.createVisualShape(
        shapeType=pb.GEOM_MESH,
        fileName=random_obj,
        meshScale=scale,
        rgbaColor=[np.random.rand(), np.random.rand(), np.random.rand(), 1]
    )
    
    return pb.createMultiBody(mass, collision_shape, visual_shape, position)

三、驱动支撑物体运动（角度控制） #

要实现支撑物体的倾斜控制，推荐用关节控制（比直接施力更稳定）：

步骤1：给支撑物体添加旋转关节 #

将支撑物体挂载到固定基座上，添加绕Y轴的旋转关节，用于控制倾斜角度：

def create_support_with_joint(self):
    # 创建固定基座（无质量，仅作为关节锚点）
    base_id = pb.createMultiBody(baseMass=0, baseCollisionShapeIndex=-1)
    # 创建随机形状的支撑物
    support_id = self.load_random_obj(mass=1.0, position=[0,0,0])
    # 创建旋转关节（绕Y轴）
    joint_id = pb.createConstraint(
        parentBodyUniqueId=base_id,
        parentLinkIndex=-1,
        childBodyUniqueId=support_id,
        childLinkIndex=-1,
        jointType=pb.JOINT_REVOLUTE,
        jointAxis=[0,1,0],
        parentFramePosition=[0,0,0.5],
        childFramePosition=[0,0,0]
    )
    return support_id, joint_id

步骤2：在Gym环境中控制关节角度 #

修改_take_action方法，根据动作调整支撑物的目标角度：

def _take_action(self, action):
    # action映射：0=减小角度, 1=保持, 2=增大角度
    angle_step = self.MOTOR_SPEED * self.TIME_STEP
    if action == 0:
        self.target_beam_angle -= angle_step
    elif action == 2:
        self.target_beam_angle += angle_step
    
    # 限制角度在设定范围内
    self.target_beam_angle = np.clip(self.target_beam_angle, self.obs_low_bounds[3], self.obs_high_bounds[3])
    
    # 应用关节位置控制
    pb.setJointMotorControl2(
        bodyUniqueId=self.support_id,
        jointIndex=0,
        controlMode=pb.POSITION_CONTROL,
        targetPosition=np.deg2rad(self.target_beam_angle),
        force=100
    )
    
    # 模拟对应时间步的物理过程
    for _ in range(int(self.TIME_STEP / 0.01)):
        pb.stepSimulation()

四、完整整合后的Gym环境 #

下面是整合了所有功能的完整环境代码，包含随机形状生成、物理控制、观测与奖励逻辑：

import os
import gym
import numpy as np
import pybullet as pb
import pybullet_data
import random
import glob
from gym import spaces
import gym.utils.seeding as seeding

class RandomShapeBalanceEnv(gym.Env):
    def __init__(self, obs_low_bounds = np.array([-0.5, -0.5, -2, -45]), obs_high_bounds = np.array([0.5, 0.5, 2, 45])):
        # 初始化PyBullet
        self.physicsClient = pb.connect(pb.GUI)
        pb.setAdditionalSearchPath(pybullet_data.getDataPath())
        pb.setGravity(0,0,-9.8)
        pb.setTimeStep(0.01)
        
        # 环境参数
        self.ACC_GRAV = 9.8 # m/s²
        self.MOTOR_SPEED = 45 # deg/s
        self.TIME_STEP = 0.1 # s
        self.obs_low_bounds = obs_low_bounds
        self.obs_high_bounds = obs_high_bounds
        
        # 动作空间：0=减小角度,1=保持,2=增大角度
        self.action_space = spaces.Discrete(3)
        # 观测空间：[小球X位置, 小球X速度, 支撑物角度, 支撑物角速度]
        self.observation_space = spaces.Box(low=self.obs_low_bounds, high=self.obs_high_bounds, dtype=np.float32)
        self.reward_range = (-1, 1)
        
        # 初始化环境
        self._seed()
        self.reset()
    
    def _seed(self, seed=None):
        self.np_random, seed = seeding.np_random(seed)
        return [seed]
    
    def create_random_convex_shape(self, mass=1.0, position=[0,0,1]):
        num_vertices = self.np_random.randint(3, 8)
        angles = np.linspace(0, 2*np.pi, num_vertices, endpoint=False)
        radii = self.np_random.uniform(0.2, 0.5, num_vertices)
        vertices = np.array([[r*np.cos(a), r*np.sin(a), 0] for r,a in zip(radii, angles)])
        
        collision_shape = pb.createCollisionShape(shapeType=pb.GEOM_MESH, vertices=vertices.tolist())
        visual_shape = pb.createVisualShape(shapeType=pb.GEOM_MESH, vertices=vertices.tolist(), rgbaColor=[self.np_random.rand(), self.np_random.rand(), self.np_random.rand(), 1])
        
        return pb.createMultiBody(mass, collision_shape, visual_shape, position)
    
    def load_random_obj(self, mass=1.0, position=[0,0,1], scale=[0.1,0.1,0.1]):
        obj_files = glob.glob("procedural_objects/**/*.obj", recursive=True)
        if not obj_files:
            return self.create_random_convex_shape(mass, position)
        
        random_obj = self.np_random.choice(obj_files)
        collision_shape = pb.createCollisionShape(shapeType=pb.GEOM_MESH, fileName=random_obj, meshScale=scale)
        visual_shape = pb.createVisualShape(shapeType=pb.GEOM_MESH, fileName=random_obj, meshScale=scale, rgbaColor=[self.np_random.rand(), self.np_random.rand(), self.np_random.rand(), 1])
        
        return pb.createMultiBody(mass, collision_shape, visual_shape, position)
    
    def create_support_with_joint(self):
        base_id = pb.createMultiBody(baseMass=0, baseCollisionShapeIndex=-1)
        support_id = self.load_random_obj(mass=1.0, position=[0,0,0])
        joint_id = pb.createConstraint(
            parentBodyUniqueId=base_id,
            parentLinkIndex=-1,
            childBodyUniqueId=support_id,
            childLinkIndex=-1,
            jointType=pb.JOINT_REVOLUTE,
            jointAxis=[0,1,0],
            parentFramePosition=[0,0,0.5],
            childFramePosition=[0,0,0]
        )
        return support_id, joint_id
    
    def reset(self):
        pb.resetSimulation()
        pb.setGravity(0,0,-9.8)
        pb.setTimeStep(0.01)
        
        # 加载地面
        self.plane_id = pb.loadURDF('plane.urdf')
        
        # 创建支撑物和关节
        self.support_id, self.joint_id = self.create_support_with_joint()
        self.target_beam_angle = 0
        
        # 创建随机形状的"小球"（用小尺寸凸多边形模拟）
        self.ball_id = self.load_random_obj(mass=0.1, position=[0,0,1.0], scale=[0.05,0.05,0.05])
        
        # 运行几步让物理状态稳定
        for _ in range(10):
            pb.stepSimulation()
        
        self._envStepCounter = 0
        return self._compute_observation()
    
    def _compute_observation(self):
        # 获取小球状态
        ball_pos, _ = pb.getBasePositionAndOrientation(self.ball_id)
        ball_lin_vel, _ = pb.getBaseVelocity(self.ball_id)
        
        # 获取支撑物角度（转成角度）
        joint_state = pb.getJointState(self.support_id, 0)
        beam_angle = np.rad2deg(joint_state[0])
        beam_ang_vel = np.rad2deg(joint_state[1])
        
        return np.array([ball_pos[0], ball_lin_vel[0], beam_angle, beam_ang_vel], dtype=np.float32)
    
    def _compute_reward(self):
        obs = self._compute_observation()
        ball_x = obs[0]
        beam_angle = obs[2]
        
        # 奖励逻辑：小球越靠近中心、支撑物角度越小，奖励越高
        reward = 1.0 - (abs(ball_x) * 2 + abs(beam_angle) / 45)
        reward = np.clip(reward, -1, 1)
        
        # 小球掉落时给予惩罚
        ball_pos, _ = pb.getBasePositionAndOrientation(self.ball_id)
        if ball_pos[2] < 0.3:
            reward = -1.0
        
        return reward
    
    def _compute_done(self):
        ball_pos, _ = pb.getBasePositionAndOrientation(self.ball_id)
        # 小球掉落或步数过多时结束回合
        return ball_pos[2] < 0.3 or self._envStepCounter >= 1000
    
    def _take_action(self, action):
        angle_step = self.MOTOR_SPEED * self.TIME_STEP
        if action == 0:
            self.target_beam_angle -= angle_step
        elif action == 2:
            self.target_beam_angle += angle_step
        
        self.target_beam_angle = np.clip(self.target_beam_angle, self.obs_low_bounds[3], self.obs_high_bounds[3])
        
        pb.setJointMotorControl2(
            bodyUniqueId=self.support_id,
            jointIndex=0,
            controlMode=pb.POSITION_CONTROL,
            targetPosition=np.deg2rad(self.target_beam_angle),
            force=100
        )
        
        # 模拟对应时间步的物理过程
        for _ in range(int(self.TIME_STEP / 0.01)):
            pb.stepSimulation()
    
    def step(self, action):
        self._envStepCounter += 1
        
        self._take_action(action)
        obs = self._compute_observation()
        reward = self._compute_reward()
        done = self._compute_done()
        
        return obs, reward, done, {}
    
    def close(self):
        pb.disconnect()

五、关键注意事项 #

1. 碰撞形状稳定性：如果加载的.obj是凹面模型，PyBullet碰撞检测可能不稳定，建议用pybullet_utils的凸分解工具处理模型。
2. 物理参数调整：可以修改关节控制的force参数、时间步长，优化物体运动的平滑度。
3. 奖励函数定制：你可以根据需求调整奖励逻辑，比如加入小球速度惩罚、目标位置追踪等。
4. 可复现性：通过_seed方法控制随机种子，确保实验结果可复现。

六、测试代码 #

if __name__ == "__main__":
    env = RandomShapeBalanceEnv()
    obs = env.reset()
    for _ in range(500):
        action = env.action_space.sample() # 随机动作，后续可替换为强化学习算法
        obs, reward, done, info = env.step(action)
        if done:
            obs = env.reset()
    env.close()

内容的提问来源于stack exchange，提问作者Emilio