mirror of
https://github.com/bulletphysics/bullet3
synced 2025-01-07 08:10:08 +00:00
496e614f77
This should get rid of the annoying warnings that pop up whenever an environment is created. Numpy's default datatype is float64, gym's is float32, the env is actually float32, but the initialization doesn't make it explicit so it causes redundant warnings.
```
/path/lib/python3.9/site-packages/gym/logger.py:34: UserWarning: WARN: Box bound precision lowered by casting to float32 warnings.warn(colorize("%s: %s" % ("WARN", msg % args), "yellow"))
```
409 lines
14 KiB
Python
409 lines
14 KiB
Python
import pybullet
|
|
import gym, gym.spaces, gym.utils
|
|
import numpy as np
|
|
import os, inspect
|
|
currentdir = os.path.dirname(os.path.abspath(inspect.getfile(inspect.currentframe())))
|
|
parentdir = os.path.dirname(currentdir)
|
|
os.sys.path.insert(0, parentdir)
|
|
import pybullet_data
|
|
|
|
|
|
class XmlBasedRobot:
|
|
"""
|
|
Base class for mujoco .xml based agents.
|
|
"""
|
|
|
|
self_collision = True
|
|
|
|
def __init__(self, robot_name, action_dim, obs_dim, self_collision):
|
|
self.parts = None
|
|
self.objects = []
|
|
self.jdict = None
|
|
self.ordered_joints = None
|
|
self.robot_body = None
|
|
|
|
high = np.ones([action_dim], dtype=np.float32)
|
|
self.action_space = gym.spaces.Box(-high, high, dtype=np.float32)
|
|
high = np.inf * np.ones([obs_dim], dtype=np.float32)
|
|
self.observation_space = gym.spaces.Box(-high, high, dtype=np.float32)
|
|
|
|
#self.model_xml = model_xml
|
|
self.robot_name = robot_name
|
|
self.self_collision = self_collision
|
|
|
|
def addToScene(self, bullet_client, bodies):
|
|
self._p = bullet_client
|
|
|
|
if self.parts is not None:
|
|
parts = self.parts
|
|
else:
|
|
parts = {}
|
|
|
|
if self.jdict is not None:
|
|
joints = self.jdict
|
|
else:
|
|
joints = {}
|
|
|
|
if self.ordered_joints is not None:
|
|
ordered_joints = self.ordered_joints
|
|
else:
|
|
ordered_joints = []
|
|
|
|
if np.isscalar(bodies): # streamline the case where bodies is actually just one body
|
|
bodies = [bodies]
|
|
|
|
dump = 0
|
|
for i in range(len(bodies)):
|
|
if self._p.getNumJoints(bodies[i]) == 0:
|
|
part_name, robot_name = self._p.getBodyInfo(bodies[i])
|
|
self.robot_name = robot_name.decode("utf8")
|
|
part_name = part_name.decode("utf8")
|
|
parts[part_name] = BodyPart(self._p, part_name, bodies, i, -1)
|
|
for j in range(self._p.getNumJoints(bodies[i])):
|
|
self._p.setJointMotorControl2(bodies[i],
|
|
j,
|
|
pybullet.POSITION_CONTROL,
|
|
positionGain=0.1,
|
|
velocityGain=0.1,
|
|
force=0)
|
|
jointInfo = self._p.getJointInfo(bodies[i], j)
|
|
joint_name = jointInfo[1]
|
|
part_name = jointInfo[12]
|
|
|
|
joint_name = joint_name.decode("utf8")
|
|
part_name = part_name.decode("utf8")
|
|
|
|
if dump: print("ROBOT PART '%s'" % part_name)
|
|
if dump:
|
|
print(
|
|
"ROBOT JOINT '%s'" % joint_name
|
|
) # limits = %+0.2f..%+0.2f effort=%0.3f speed=%0.3f" % ((joint_name,) + j.limits()) )
|
|
|
|
parts[part_name] = BodyPart(self._p, part_name, bodies, i, j)
|
|
|
|
if part_name == self.robot_name:
|
|
self.robot_body = parts[part_name]
|
|
|
|
if i == 0 and j == 0 and self.robot_body is None: # if nothing else works, we take this as robot_body
|
|
parts[self.robot_name] = BodyPart(self._p, self.robot_name, bodies, 0, -1)
|
|
self.robot_body = parts[self.robot_name]
|
|
|
|
if joint_name[:6] == "ignore":
|
|
Joint(self._p, joint_name, bodies, i, j).disable_motor()
|
|
continue
|
|
|
|
if joint_name[:8] != "jointfix":
|
|
joints[joint_name] = Joint(self._p, joint_name, bodies, i, j)
|
|
ordered_joints.append(joints[joint_name])
|
|
|
|
joints[joint_name].power_coef = 100.0
|
|
|
|
# TODO: Maybe we need this
|
|
# joints[joint_name].power_coef, joints[joint_name].max_velocity = joints[joint_name].limits()[2:4]
|
|
# self.ordered_joints.append(joints[joint_name])
|
|
# self.jdict[joint_name] = joints[joint_name]
|
|
|
|
return parts, joints, ordered_joints, self.robot_body
|
|
|
|
def reset_pose(self, position, orientation):
|
|
self.parts[self.robot_name].reset_pose(position, orientation)
|
|
|
|
|
|
class MJCFBasedRobot(XmlBasedRobot):
|
|
"""
|
|
Base class for mujoco .xml based agents.
|
|
"""
|
|
|
|
def __init__(self, model_xml, robot_name, action_dim, obs_dim, self_collision=True):
|
|
XmlBasedRobot.__init__(self, robot_name, action_dim, obs_dim, self_collision)
|
|
self.model_xml = model_xml
|
|
self.doneLoading = 0
|
|
|
|
def reset(self, bullet_client):
|
|
|
|
self._p = bullet_client
|
|
#print("Created bullet_client with id=", self._p._client)
|
|
if (self.doneLoading == 0):
|
|
self.ordered_joints = []
|
|
self.doneLoading = 1
|
|
if self.self_collision:
|
|
self.objects = self._p.loadMJCF(os.path.join(pybullet_data.getDataPath(), "mjcf",
|
|
self.model_xml),
|
|
flags=pybullet.URDF_USE_SELF_COLLISION |
|
|
pybullet.URDF_USE_SELF_COLLISION_EXCLUDE_ALL_PARENTS |
|
|
pybullet.URDF_GOOGLEY_UNDEFINED_COLORS )
|
|
self.parts, self.jdict, self.ordered_joints, self.robot_body = self.addToScene(
|
|
self._p, self.objects)
|
|
else:
|
|
self.objects = self._p.loadMJCF(
|
|
os.path.join(pybullet_data.getDataPath(), "mjcf", self.model_xml, flags = pybullet.URDF_GOOGLEY_UNDEFINED_COLORS))
|
|
self.parts, self.jdict, self.ordered_joints, self.robot_body = self.addToScene(
|
|
self._p, self.objects)
|
|
self.robot_specific_reset(self._p)
|
|
|
|
s = self.calc_state(
|
|
) # optimization: calc_state() can calculate something in self.* for calc_potential() to use
|
|
|
|
return s
|
|
|
|
def calc_potential(self):
|
|
return 0
|
|
|
|
|
|
class URDFBasedRobot(XmlBasedRobot):
|
|
"""
|
|
Base class for URDF .xml based robots.
|
|
"""
|
|
|
|
def __init__(self,
|
|
model_urdf,
|
|
robot_name,
|
|
action_dim,
|
|
obs_dim,
|
|
basePosition=[0, 0, 0],
|
|
baseOrientation=[0, 0, 0, 1],
|
|
fixed_base=False,
|
|
self_collision=False):
|
|
XmlBasedRobot.__init__(self, robot_name, action_dim, obs_dim, self_collision)
|
|
|
|
self.model_urdf = model_urdf
|
|
self.basePosition = basePosition
|
|
self.baseOrientation = baseOrientation
|
|
self.fixed_base = fixed_base
|
|
|
|
def reset(self, bullet_client):
|
|
self._p = bullet_client
|
|
self.ordered_joints = []
|
|
|
|
print(os.path.join(os.path.dirname(__file__), "data", self.model_urdf))
|
|
|
|
if self.self_collision:
|
|
self.parts, self.jdict, self.ordered_joints, self.robot_body = self.addToScene(
|
|
self._p,
|
|
self._p.loadURDF(os.path.join(pybullet_data.getDataPath(), self.model_urdf),
|
|
basePosition=self.basePosition,
|
|
baseOrientation=self.baseOrientation,
|
|
useFixedBase=self.fixed_base,
|
|
flags=pybullet.URDF_USE_SELF_COLLISION | pybullet.URDF_GOOGLEY_UNDEFINED_COLORS))
|
|
else:
|
|
self.parts, self.jdict, self.ordered_joints, self.robot_body = self.addToScene(
|
|
self._p,
|
|
self._p.loadURDF(os.path.join(pybullet_data.getDataPath(), self.model_urdf),
|
|
basePosition=self.basePosition,
|
|
baseOrientation=self.baseOrientation,
|
|
useFixedBase=self.fixed_base, flags = pybullet.URDF_GOOGLEY_UNDEFINED_COLORS))
|
|
|
|
self.robot_specific_reset(self._p)
|
|
|
|
s = self.calc_state(
|
|
) # optimization: calc_state() can calculate something in self.* for calc_potential() to use
|
|
self.potential = self.calc_potential()
|
|
|
|
return s
|
|
|
|
def calc_potential(self):
|
|
return 0
|
|
|
|
|
|
class SDFBasedRobot(XmlBasedRobot):
|
|
"""
|
|
Base class for SDF robots in a Scene.
|
|
"""
|
|
|
|
def __init__(self,
|
|
model_sdf,
|
|
robot_name,
|
|
action_dim,
|
|
obs_dim,
|
|
basePosition=[0, 0, 0],
|
|
baseOrientation=[0, 0, 0, 1],
|
|
fixed_base=False,
|
|
self_collision=False):
|
|
XmlBasedRobot.__init__(self, robot_name, action_dim, obs_dim, self_collision)
|
|
|
|
self.model_sdf = model_sdf
|
|
self.fixed_base = fixed_base
|
|
|
|
def reset(self, bullet_client):
|
|
self._p = bullet_client
|
|
|
|
self.ordered_joints = []
|
|
|
|
self.parts, self.jdict, self.ordered_joints, self.robot_body = self.addToScene(
|
|
self._p, # TODO: Not sure if this works, try it with kuka
|
|
self._p.loadSDF(os.path.join("models_robot", self.model_sdf)))
|
|
|
|
self.robot_specific_reset(self._p)
|
|
|
|
s = self.calc_state(
|
|
) # optimization: calc_state() can calculate something in self.* for calc_potential() to use
|
|
self.potential = self.calc_potential()
|
|
|
|
return s
|
|
|
|
def calc_potential(self):
|
|
return 0
|
|
|
|
|
|
class Pose_Helper: # dummy class to comply to original interface
|
|
|
|
def __init__(self, body_part):
|
|
self.body_part = body_part
|
|
|
|
def xyz(self):
|
|
return self.body_part.current_position()
|
|
|
|
def rpy(self):
|
|
return pybullet.getEulerFromQuaternion(self.body_part.current_orientation())
|
|
|
|
def orientation(self):
|
|
return self.body_part.current_orientation()
|
|
|
|
|
|
class BodyPart:
|
|
|
|
def __init__(self, bullet_client, body_name, bodies, bodyIndex, bodyPartIndex):
|
|
self.bodies = bodies
|
|
self._p = bullet_client
|
|
self.bodyIndex = bodyIndex
|
|
self.bodyPartIndex = bodyPartIndex
|
|
self.initialPosition = self.current_position()
|
|
self.initialOrientation = self.current_orientation()
|
|
self.bp_pose = Pose_Helper(self)
|
|
|
|
def state_fields_of_pose_of(
|
|
self, body_id,
|
|
link_id=-1): # a method you will most probably need a lot to get pose and orientation
|
|
if link_id == -1:
|
|
(x, y, z), (a, b, c, d) = self._p.getBasePositionAndOrientation(body_id)
|
|
else:
|
|
(x, y, z), (a, b, c, d), _, _, _, _ = self._p.getLinkState(body_id, link_id)
|
|
return np.array([x, y, z, a, b, c, d])
|
|
|
|
def get_position(self):
|
|
return self.current_position()
|
|
|
|
def get_pose(self):
|
|
return self.state_fields_of_pose_of(self.bodies[self.bodyIndex], self.bodyPartIndex)
|
|
|
|
def speed(self):
|
|
if self.bodyPartIndex == -1:
|
|
(vx, vy, vz), _ = self._p.getBaseVelocity(self.bodies[self.bodyIndex])
|
|
else:
|
|
(x, y, z), (a, b, c, d), _, _, _, _, (vx, vy, vz), (vr, vp, vy) = self._p.getLinkState(
|
|
self.bodies[self.bodyIndex], self.bodyPartIndex, computeLinkVelocity=1)
|
|
return np.array([vx, vy, vz])
|
|
|
|
def current_position(self):
|
|
return self.get_pose()[:3]
|
|
|
|
def current_orientation(self):
|
|
return self.get_pose()[3:]
|
|
|
|
def get_orientation(self):
|
|
return self.current_orientation()
|
|
|
|
def reset_position(self, position):
|
|
self._p.resetBasePositionAndOrientation(self.bodies[self.bodyIndex], position,
|
|
self.get_orientation())
|
|
|
|
def reset_orientation(self, orientation):
|
|
self._p.resetBasePositionAndOrientation(self.bodies[self.bodyIndex], self.get_position(),
|
|
orientation)
|
|
|
|
def reset_velocity(self, linearVelocity=[0, 0, 0], angularVelocity=[0, 0, 0]):
|
|
self._p.resetBaseVelocity(self.bodies[self.bodyIndex], linearVelocity, angularVelocity)
|
|
|
|
def reset_pose(self, position, orientation):
|
|
self._p.resetBasePositionAndOrientation(self.bodies[self.bodyIndex], position, orientation)
|
|
|
|
def pose(self):
|
|
return self.bp_pose
|
|
|
|
def contact_list(self):
|
|
return self._p.getContactPoints(self.bodies[self.bodyIndex], -1, self.bodyPartIndex, -1)
|
|
|
|
|
|
class Joint:
|
|
|
|
def __init__(self, bullet_client, joint_name, bodies, bodyIndex, jointIndex):
|
|
self.bodies = bodies
|
|
self._p = bullet_client
|
|
self.bodyIndex = bodyIndex
|
|
self.jointIndex = jointIndex
|
|
self.joint_name = joint_name
|
|
|
|
jointInfo = self._p.getJointInfo(self.bodies[self.bodyIndex], self.jointIndex)
|
|
self.lowerLimit = jointInfo[8]
|
|
self.upperLimit = jointInfo[9]
|
|
|
|
self.power_coeff = 0
|
|
|
|
def set_state(self, x, vx):
|
|
self._p.resetJointState(self.bodies[self.bodyIndex], self.jointIndex, x, vx)
|
|
|
|
def current_position(self): # just some synonyme method
|
|
return self.get_state()
|
|
|
|
def current_relative_position(self):
|
|
pos, vel = self.get_state()
|
|
pos_mid = 0.5 * (self.lowerLimit + self.upperLimit)
|
|
return (2 * (pos - pos_mid) / (self.upperLimit - self.lowerLimit), 0.1 * vel)
|
|
|
|
def get_state(self):
|
|
x, vx, _, _ = self._p.getJointState(self.bodies[self.bodyIndex], self.jointIndex)
|
|
return x, vx
|
|
|
|
def get_position(self):
|
|
x, _ = self.get_state()
|
|
return x
|
|
|
|
def get_orientation(self):
|
|
_, r = self.get_state()
|
|
return r
|
|
|
|
def get_velocity(self):
|
|
_, vx = self.get_state()
|
|
return vx
|
|
|
|
def set_position(self, position):
|
|
self._p.setJointMotorControl2(self.bodies[self.bodyIndex],
|
|
self.jointIndex,
|
|
pybullet.POSITION_CONTROL,
|
|
targetPosition=position)
|
|
|
|
def set_velocity(self, velocity):
|
|
self._p.setJointMotorControl2(self.bodies[self.bodyIndex],
|
|
self.jointIndex,
|
|
pybullet.VELOCITY_CONTROL,
|
|
targetVelocity=velocity)
|
|
|
|
def set_motor_torque(self, torque): # just some synonyme method
|
|
self.set_torque(torque)
|
|
|
|
def set_torque(self, torque):
|
|
self._p.setJointMotorControl2(bodyIndex=self.bodies[self.bodyIndex],
|
|
jointIndex=self.jointIndex,
|
|
controlMode=pybullet.TORQUE_CONTROL,
|
|
force=torque) #, positionGain=0.1, velocityGain=0.1)
|
|
|
|
def reset_current_position(self, position, velocity): # just some synonyme method
|
|
self.reset_position(position, velocity)
|
|
|
|
def reset_position(self, position, velocity):
|
|
self._p.resetJointState(self.bodies[self.bodyIndex],
|
|
self.jointIndex,
|
|
targetValue=position,
|
|
targetVelocity=velocity)
|
|
self.disable_motor()
|
|
|
|
def disable_motor(self):
|
|
self._p.setJointMotorControl2(self.bodies[self.bodyIndex],
|
|
self.jointIndex,
|
|
controlMode=pybullet.POSITION_CONTROL,
|
|
targetPosition=0,
|
|
targetVelocity=0,
|
|
positionGain=0.1,
|
|
velocityGain=0.1,
|
|
force=0)
|