def make_robot_xml(cls, xml, robot_control_params):
if robot_control_params.is_joint_actuated():
# Modifying xml is required because setting eq_active only was not enough to fully
# disable the mocap weld constraint. In my tests, setting eq_active to false would
# disable the constraint, but somehow the arm would not move when the joints were
# commanded. Removing from xml here seems to have the right effect.
xml.remove_objects_by_name("mocap_weld")
# Also add the actuations that are removed in the xml by default (since TCP does
# not need them).
joint_subdir = robot_control_params.arm_joint_calibration_path
xml.append(
MujocoXML.parse(
f"robot/ur16e/jointspec/calibrations/{joint_subdir}/ur16e_ik_class.xml"
))
xml.append(
MujocoXML.parse(
f"robot/ur16e/jointspec/calibrations/{joint_subdir}/joint_actuations.xml"
))
else:
# If not joint control mode or ik solver mode, use mocap defaults for joint parameters
xml.append(
MujocoXML.parse("robot/ur16e/jointspec/ur16e_mocap_class.xml"))
# Add gripper actuators now (after joint actuators if required).
xml.append(MujocoXML.parse("robot/ur16e/gripper_actuators.xml"))
return xml
def test_set_attributes_mixed_precision():
main = (MujocoXML().add_default_compiler_directive().append(
MujocoXML.from_string(XML_BALL).set_named_objects_attr(
"ball", pos=[1, 1e-8, 1e-12])))
simulation = SimulationInterface(main.build())
ball_id = simulation.sim.model.body_name2id("ball")
ball_pos = simulation.sim.model.body_pos[ball_id]
target_pos = np.array([1, 1e-8, 1e-12])
# test relative error cause absolute error can be quite small either way
assert np.linalg.norm((ball_pos / target_pos) - 1) < 1e-6
def make_mesh_object(name: str, files: List[str], scale: float) -> MujocoXML:
# Center mesh properly by offsetting with center position of combined mesh.
mesh = get_combined_mesh(files)
pos = -mesh.center_mass * scale
pos_string = " ".join(map(str, pos))
scale_string = " ".join(map(str, [scale] * 3))
assets = [
f'<mesh file="{file}" name="{name}-{idx}" scale="{scale_string}" />'
for idx, file in enumerate(files)
]
geoms = [
f'<geom type="mesh" mesh="{name}-{idx}" pos="{pos_string}"/>'
for idx in range(len(files))
]
assets_xml = "\n".join(assets)
geoms_xml = "\n".join(geoms)
xml_source = f"""
<mujoco>
<asset>
{assets_xml}
</asset>
<worldbody>
<body name="{name}" pos="0.0 0.0 0.0">
{geoms_xml}
<joint name="{name}:joint" type="free"/>
</body>
</worldbody>
</mujoco>
"""
return MujocoXML.from_string(xml_source)
def test_forward_kinematics_on_inverted_pendulum():
mxml = MujocoXML.parse(
"test/inverted_pendulum/inverted_double_pendulum.xml").add_name_prefix(
"ivp:")
simulation = SimulationInterface(mxml.build())
simulation.register_joint_group("pendulum", "ivp:")
joint_names = list(map(lambda x: "ivp:%s" % x, ["hinge", "hinge2"]))
site_names = list(map(lambda x: "ivp:%s" % x, ["hinge2_site", "tip"]))
KIN = ForwardKinematics.prepare(mxml, "ivp:cart", np.zeros(3), np.zeros(3),
site_names, joint_names)
for _ in range(5):
simulation.mj_sim.data.ctrl[0] = random.random()
for _ in range(100):
simulation.step()
simulation.forward()
site_positions = np.array([
simulation.mj_sim.data.get_site_xpos(site) for site in site_names
])
joint_pos = simulation.get_qpos("pendulum")
kinemetics_positions = KIN.compute(joint_pos, return_joint_pos=True)
assert (np.abs(site_positions - kinemetics_positions[:2]) < 1e-6).all()
assert (np.abs(site_positions[0] - kinemetics_positions[-1]) <
1e-6).all()
def make_objects_xml(cls, xml, simulation_params):
def get_cube_mesh(c) -> str:
return f"""<mesh name="object:letter_{c}" file="{ASSETS_DIR}/stls/openai_cube/{c}.stl" />"""
def get_cube_body(idx, c) -> str:
color = "0.71 0.61 0.3"
letter_material = "cube:letter_white"
return f"""
<mujoco>
<worldbody>
<body name="object{idx}">
<joint name="object{idx}:joint" type="free"/>
<geom name="object{idx}" size="0.0285 0.0285 0.0285" type="box" rgba="{color} 1" />
<body name="object:letter_{c}:face_top" pos="0 0 0.0285">
<geom name="object:letter_{c}" pos="0 0 -0.0009" quat="0.499998 0.5 -0.500002 0.5" type="mesh" material="{letter_material}" mesh="object:letter_{c}" />
</body>
</body>
</worldbody>
</mujoco>
"""
letter_meshes = "\n".join([get_cube_mesh(c) for c in "OPENAI"])
assets_xml = f"""
<mujoco>
<asset>
<material name="cube:top_background" specular="1" shininess="0.3" />
<material name="cube:letter" specular="1" shininess="0.3" rgba="0. 0. 1 1"/>
<material name="cube:letter_white" specular="1" shininess="0.3" rgba="1 1 1 1"/>
{letter_meshes}
</asset>
</mujoco>
"""
xml.append(MujocoXML.from_string(assets_xml))
xmls = []
for idx, c in enumerate("OPENAI"):
if idx >= simulation_params.num_objects:
break
object_xml = MujocoXML.from_string(get_cube_body(idx, c))
xmls.append((object_xml, make_target(object_xml)))
return xmls
def _get_default_xml():
xml_source = """
<mujoco>
<asset>
<material name="block_mat" specular="0" shininess="0.5" reflectance="0" rgba="1 0 0 1"></material>
</asset>
</mujoco>
"""
xml = MujocoXML.from_string(xml_source)
return xml
def _make_object_xml(cls, object_id: int, config: ObjectConfig):
object_name = f"object{object_id}"
xml = (MujocoXML.parse(config.xml_path).set_objects_attr(
tag="joint", name="joint").add_name_prefix(
f"{object_name}:",
exclude_attribs=["class"]).set_objects_attr(tag="body",
name=object_name))
xml.set_objects_attrs(config.tag_args)
xml.set_objects_attrs(config.material_args)
return xml
def _build_mujoco_cube_xml(cls, xml, cube_xml_path):
xml.append(
MujocoXML.parse(cube_xml_path).remove_objects_by_name(
"annotation:outer_bound").add_name_prefix(
"cube:").set_named_objects_attr(
"cube:middle", tag="body",
pos=[1.0, 0.87,
0.2]).set_named_objects_attr("cube:middle",
tag="geom",
density=421.0))
# Target cube
xml.append(
MujocoXML.parse(cube_xml_path).remove_objects_by_name(
"annotation:outer_bound").add_name_prefix(
"target:").set_named_objects_attr("target:middle",
tag="body",
pos=[1.0, 0.87, 0.2])
# Disable collisions
.set_objects_attr(tag="geom",
group="2",
conaffinity="0",
contype="0"))
def test_remove_elem():
ball_without_joint = MujocoXML.from_string(XML_BALL).remove_objects_by_tag(
"freejoint")
ref_xml = """
<mujoco>
<worldbody>
<body name="ball">
<geom name="sphere" pos="0.00 0.00 0.00" size="0.1 0.1 0.1" type="sphere" />
</body>
</worldbody>
</mujoco>
"""
assert ref_xml.strip() == ball_without_joint.xml_string().strip()
def test_simple_mujoco_setup():
ball_one = (MujocoXML.from_string(XML_BALL).add_name_prefix(
"ball_one:").set_named_objects_attr("ball_one:ball", pos=[1, 0, 0]))
ball_two = (MujocoXML.from_string(XML_BALL).add_name_prefix(
"ball_two:").set_named_objects_attr("ball_two:ball", pos=[-1, 0, 0]))
main = (MujocoXML().add_default_compiler_directive().append(
ball_one).append(ball_two))
simulation = SimulationInterface(main.build())
simulation.register_joint_group("ball_one", "ball_one:ball_joint")
simulation.register_joint_group("ball_two", "ball_two:ball_joint")
assert simulation.get_qpos("ball_one").shape == (7, )
assert simulation.get_qpos("ball_two").shape == (7, )
assert simulation.get_qvel("ball_one").shape == (6, )
assert simulation.get_qvel("ball_two").shape == (6, )
qpos1 = np.random.randn(3)
qrot1 = uniform_quat(global_randstate)
qpos1_combined = np.concatenate([qpos1, qrot1])
qpos2 = np.random.randn(3)
qrot2 = uniform_quat(global_randstate)
qpos2_combined = np.concatenate([qpos2, qrot2])
simulation.set_qpos("ball_one", qpos1_combined)
simulation.set_qpos("ball_two", qpos2_combined)
assert np.linalg.norm(simulation.get_qpos("ball_one") -
qpos1_combined) < 1e-6
assert np.linalg.norm(simulation.get_qpos("ball_two") -
qpos2_combined) < 1e-6
def _build_mujoco_cube_xml(cls, xml, cube_xml_path):
xml.append(
MujocoXML.parse(cube_xml_path).add_name_prefix(
"cube:").set_named_objects_attr("cube:middle",
tag="body",
pos=[1.0, 0.87, 0.2])
# Delete springs for now
.remove_objects_by_prefix(prefix="cube:cubelet:spring:",
tag="joint"))
# Target
xml.append(
MujocoXML.parse(cube_xml_path).add_name_prefix(
"target:").set_named_objects_attr("target:middle",
tag="body",
pos=[1.0, 0.87, 0.2])
# Delete springs for now
.remove_objects_by_prefix(prefix="target:cubelet:spring:",
tag="joint")
# Disable collisions
.set_objects_attr(tag="geom",
group="2",
conaffinity="0",
contype="0"))
def test_more_complex_mujoco_setup():
xml = (MujocoXML().add_default_compiler_directive().append(
MujocoXML.from_string(XML_ARM).add_name_prefix(
"arm_one:").set_named_objects_attr(
"arm_one:ball", pos=[0, 1, 0])).append(
MujocoXML.from_string(XML_ARM).add_name_prefix(
"arm_two:").set_named_objects_attr("arm_two:ball",
pos=[0, -1, 0])))
simulation = SimulationInterface(xml.build())
simulation.register_joint_group("arm_one", "arm_one:")
simulation.register_joint_group("arm_one_hinge", "arm_one:hinge_joint")
simulation.register_joint_group("arm_two", "arm_two:")
simulation.register_joint_group("arm_two_hinge", "arm_two:hinge_joint")
assert simulation.get_qpos("arm_one").shape == (2, )
assert simulation.get_qvel("arm_one").shape == (2, )
assert simulation.get_qpos("arm_two").shape == (2, )
assert simulation.get_qvel("arm_two").shape == (2, )
assert simulation.get_qpos("arm_one_hinge").shape == (1, )
assert simulation.get_qvel("arm_one_hinge").shape == (1, )
assert simulation.get_qpos("arm_two_hinge").shape == (1, )
assert simulation.get_qvel("arm_two_hinge").shape == (1, )
initial_qpos_one = simulation.get_qpos("arm_one")
initial_qpos_two = simulation.get_qpos("arm_two")
simulation.set_qpos("arm_one_hinge", 0.1)
# Chech that we are setting the right hinge joint
assert np.linalg.norm(simulation.get_qpos("arm_one") -
initial_qpos_one) > 0.09
assert np.linalg.norm(simulation.get_qpos("arm_two") -
initial_qpos_two) < 1e-6
def make_world_xml(
cls, *, mujoco_timestep: float, simulation_params=None, **kwargs
):
if simulation_params is None:
simulation_params = MeshRearrangeSimParameters()
xml = super().make_world_xml(
simulation_params=simulation_params,
mujoco_timestep=mujoco_timestep,
**kwargs
)
xml.append(
MujocoXML.from_string(
CHESSBOARD_XML.format(n_obj=simulation_params.num_objects // 2)
)
)
return xml
def make_block(name: str, object_size: np.ndarray) -> MujocoXML:
"""Creates a block.
:param name: The name of the block
:param object_size: The size of the block (3-dimensional). This is half-size as per Mujoco
convention
"""
xml_source = f"""
<mujoco>
<worldbody>
<body name="{name}" pos="0.0 0.0 0.0">
<geom type="box" rgba="0.0 0.0 0.0 0.0" material="block_mat"/>
<joint name="{name}:joint" type="free"/>
</body>
</worldbody>
</mujoco>
"""
xml = MujocoXML.from_string(xml_source).set_objects_attr(
tag="geom", size=object_size
)
return xml
def make_world_xml(cls,
*,
mujoco_timestep: float,
simulation_params=None,
**kwargs):
if simulation_params is None:
simulation_params = HoldoutRearrangeSimParameters()
world_xml = super().make_world_xml(simulation_params=simulation_params,
mujoco_timestep=mujoco_timestep)
table_dims = cls.get_table_dimensions_from_xml(world_xml)
for object_id, object_config in enumerate(
simulation_params.scene_object_configs):
world_xml.append(
cls._make_scene_object_xml(
object_id,
object_config,
table_dims,
))
if simulation_params.shared_settings:
world_xml.append(MujocoXML.parse(
simulation_params.shared_settings))
return world_xml
def build(cls, n_substeps: int = 10):
"""Construct a ShadowHandReachSimulation object.
:param n_substeps: (int) sim.nsubsteps, num of substeps
:return: a ShadowHandReachSimulation object with properly constructed sim.
"""
xml = MujocoXML()
xml.add_default_compiler_directive()
xml.append(
MujocoXML.parse(cls.FLOOR_XML).set_named_objects_attr(
"floor", tag="body", pos=[1, 1, 0]))
target = MujocoXML.parse(cls.TARGET_XML)
colors = [
[1.0, 0.0, 0.0, 1.0],
[0.0, 1.0, 0.0, 1.0],
[0.0, 0.0, 1.0, 1.0],
[1.0, 1.0, 0.0, 1.0],
[1.0, 0.0, 1.0, 1.0],
]
for site, color in zip(FINGERTIP_SITE_NAMES, colors):
target.set_named_objects_attr(
f"target_{site}",
pos=[0.5, 0.5, 0.0],
type="sphere",
rgba=color,
size=0.005,
)
xml.append(target)
xml.append(
MujocoXML.parse(cls.HAND_XML).add_name_prefix(
"robot0:").set_named_objects_attr(
"robot0:hand_mount",
tag="body",
pos=[1.0, 1.25, 0.15],
euler=[np.pi / 2, 0, np.pi],
).remove_objects_by_name("robot0:annotation:outer_bound")
# Remove hand base free joint so that hand is immovable
.remove_objects_by_name("robot0:hand_base"))
xml.append(MujocoXML.parse(cls.LIGHT_XML))
simulation = cls(xml.build(nsubsteps=n_substeps))
# Move fingers out of the way.
simulation.shadow_hand.set_position_control(
simulation.shadow_hand.denormalize_position_control(
simulation.shadow_hand.zero_control()))
for _ in range(20):
simulation.step()
return simulation
def build(
cls, n_substeps: int = 10, timestep: float = 0.008, name_prefix: str = "robot0:"
):
"""
Construct MjSim object for this simulation
:param name_prefix - to append to names of all objects in the MuJoCo model of the hand;
by default no prefix is appended.
"""
xml = MujocoXML()
xml.add_default_compiler_directive()
max_contacts_params = dict(njmax=2000, nconmax=200)
xml.append(
MujocoXML.parse(cls.FLOOR_XML).set_named_objects_attr(
"floor", tag="body", pos=[1, 1, 0]
)
)
xml.append(
MujocoXML.parse(cls.HAND_XML)
.add_name_prefix(name_prefix)
.set_objects_attr(tag="size", **max_contacts_params)
.set_objects_attr(tag="option", timestep=timestep)
.set_named_objects_attr(
f"{name_prefix}hand_mount",
tag="body",
pos=[1.0, 1.25, 0.15],
euler=[np.pi / 2, 0, np.pi],
)
.remove_objects_by_name(f"{name_prefix}annotation:outer_bound")
# Remove hand base free joint so that hand is immovable
.remove_objects_by_name(f"{name_prefix}hand_base")
)
xml.append(MujocoXML.parse(cls.LIGHT_XML))
return cls(sim=xml.build(nsubsteps=n_substeps), hand_prefix=name_prefix)
def _build_mujoco_cube_xml(cls, xml, cube_xml_path):
xml.append(
MujocoXML.parse(cube_xml_path).add_name_prefix(
"cube:").set_named_objects_attr("cube:middle",
tag="body",
pos=[1.0, 0.87, 0.2])
# Leave +/- z driver joints
.remove_objects_by_name(
names="cube:cubelet:driver:neg_x",
tag="joint").remove_objects_by_name(
names="cube:cubelet:driver:pos_x",
tag="joint").remove_objects_by_name(
names="cube:cubelet:driver:neg_y",
tag="joint").remove_objects_by_name(
names="cube:cubelet:driver:pos_y", tag="joint")
# Remove x/y cubelet hinge joints
.remove_objects_by_prefix(prefix="cube:cubelet:rotx:",
tag="joint").remove_objects_by_prefix(
prefix="cube:cubelet:roty:",
tag="joint")
# Delete springs for now
.remove_objects_by_prefix(prefix="cube:cubelet:spring:",
tag="joint")
# Remove remaining cubelet joints we're not interested in
.remove_objects_by_name(
names=[
"cube:cubelet:rotz:neg_x_pos_y",
"cube:cubelet:rotz:neg_x_neg_y",
"cube:cubelet:rotz:pos_x_pos_y",
"cube:cubelet:rotz:pos_x_neg_y",
],
tag="joint",
))
# Target
xml.append(
MujocoXML.parse(cube_xml_path).add_name_prefix(
"target:").set_named_objects_attr("target:middle",
tag="body",
pos=[1.0, 0.87, 0.2])
# Disable collisions
.set_objects_attr(tag="geom",
group="2",
conaffinity="0",
contype="0")
# Leave +/- z driver joints
.remove_objects_by_name(
names="target:cubelet:driver:neg_x",
tag="joint").remove_objects_by_name(
names="target:cubelet:driver:pos_x",
tag="joint").remove_objects_by_name(
names="target:cubelet:driver:neg_y",
tag="joint").remove_objects_by_name(
names="target:cubelet:driver:pos_y", tag="joint")
# Remove x/y cubelet hinge joints
.remove_objects_by_prefix(
prefix="target:cubelet:rotx:",
tag="joint").remove_objects_by_prefix(
prefix="target:cubelet:roty:",
tag="joint").remove_objects_by_prefix(
prefix="target:cubelet:spring:", tag="joint")
# Remove remaining cubelet joints we're not interested in
.remove_objects_by_name(
names=[
"target:cubelet:rotz:neg_x_pos_y",
"target:cubelet:rotz:neg_x_neg_y",
"target:cubelet:rotz:pos_x_pos_y",
"target:cubelet:rotz:pos_x_neg_y",
],
tag="joint",
))
def build(
cls,
n_substeps: int = 10,
simulation_params: CubeSimulationParameters = CubeSimulationParameters(),
):
"""
Construct a CubeSimulationInterface object with perpendicular cube.
"""
cube_xml_path, size_params = cls._get_model_xml_path_and_params(
simulation_params.cube_appearance
)
xml = MujocoXML()
xml.add_default_compiler_directive()
cls._build_mujoco_cube_xml(xml, cube_xml_path)
xml.append(
MujocoXML.parse(cls.FLOOR_XML).set_named_objects_attr(
"floor", tag="body", pos=[1, 1, 0]
)
)
xml.append(
MujocoXML.parse(cls.HAND_XML)
.add_name_prefix("robot0:")
.set_objects_attr(tag="size", **size_params)
.set_named_objects_attr(
"robot0:hand_mount",
tag="body",
pos=[1.0, 1.25, 0.15],
euler=[np.pi / 2, 0, np.pi],
)
.remove_objects_by_name("robot0:annotation:outer_bound")
# Remove hand base free joint so that hand is immovable
.remove_objects_by_name("robot0:hand_base")
)
xml.append(MujocoXML.parse(cls.LIGHT_XML))
simulation = cls(xml.build(nsubsteps=n_substeps))
if simulation_params.hide_target:
simulation.hide_target()
return simulation
FINGERTIP_SITE_NAMES = [
"S_fftip",
"S_mftip",
"S_rftip",
"S_lftip",
"S_thtip",
]
REFERENCE_SITE_NAMES = [
"phasespace_ref0",
"phasespace_ref1",
"phasespace_ref2",
]
HAND_KINEMATICS = ForwardKinematics.prepare(
MujocoXML.parse("robot/shadowhand/main.xml"),
"hand_mount",
[1.0, 1.25, 0.15],
[np.pi / 2, 0, np.pi],
REFERENCE_SITE_NAMES + FINGERTIP_SITE_NAMES,
hand_interface.JOINTS,
)
ZERO_JOINT_POSITIONS = HAND_KINEMATICS.compute(
np.zeros(len(hand_interface.JOINTS)))
REFERENCE_POSITIONS = ZERO_JOINT_POSITIONS[:3]
def hand_forward_kinematics(qpos, return_joint_pos=False):
""" Calculate forward kinematics of the hand """
return HAND_KINEMATICS.compute(qpos, return_joint_pos)[3:]
def make_composed_mesh_object(
name: str,
primitives: List[str],
random_state: np.random.RandomState,
mesh_size_range: tuple = (0.01, 0.1),
attachment: str = "random",
object_size: Optional[float] = None,
) -> MujocoXML:
"""
Composes an object out of mesh primitives by combining them in a random but systematic
way. In the resulting object, all meshes are guaranteed to be connected.
:param name: The name of the resulting object.
:param primitives: A list of STL files that will be used as primitives in the provided order.
:param random_state: The random state used for sampling.
:param mesh_size_range: Each mesh is randomly resized (iid per dimension) but each side is
guaranteed to be within this size. This is full-size, not half-size.
:param attachment: How primitives are connected. If "random", the parent geom is randomly
selected from the already placed geoms. If "last", the geom that was place last is used.
:param object_size: If this is not None, the final object) will be re-scaled so that the longest
side has exactly object_size half-size. This parameter is in half-size, as per Mujoco
convention.
:return: a MujocoXML object.
"""
assert 0 <= mesh_size_range[0] <= mesh_size_range[1]
assert attachment in ["random", "last"]
def compute_pos_and_size(geoms):
min_max_xyzs = np.array([geom.min_max_xyz() for geom in geoms])
min_xyz = np.min(min_max_xyzs[:, 0, :], axis=0)
max_xyz = np.max(min_max_xyzs[:, 1, :], axis=0)
size = (max_xyz - min_xyz) / 2.0
pos = min_xyz + size
return pos, size
geoms: List[MeshGeom] = []
for i, mesh_path in enumerate(primitives):
# Load mesh.
mesh = trimesh.load(mesh_path)
# Scale randomly but such that the mesh is within mesh_size_range.
min_scale = mesh_size_range[0] / mesh.bounding_box.extents
max_scale = mesh_size_range[1] / mesh.bounding_box.extents
assert min_scale.shape == max_scale.shape == (3,)
scale = random_state.uniform(min_scale, max_scale) * random_state.choice(
[-1, 1], size=3
)
assert scale.shape == (3,)
scale_matrix = np.eye(4)
np.fill_diagonal(scale_matrix[:3, :3], scale)
# Rotate randomly.
quat = uniform_quat(random_state)
rotation_matrix = np.eye(4)
rotation_matrix[:3, :3] = quat2mat(quat)
# Apply transformations. Apply scaling first since we computed the scale
# in the original reference frame! In principle, we could also sheer the
# object, but we currently do not do this.
mesh.apply_transform(scale_matrix)
mesh.apply_transform(rotation_matrix)
if len(geoms) == 0:
pos = -mesh.center_mass
else:
if attachment == "random":
parent_geom = random_state.choice(geoms)
elif attachment == "last":
parent_geom = geoms[-1]
else:
raise ValueError()
# We sample 10 points here because sample_surface sometimes returns less points
# than we requested (unclear why).
surface_pos = trimesh.sample.sample_surface(parent_geom.mesh, 10)[0][0]
pos = parent_geom.pos + (surface_pos - mesh.center_mass)
geom = MeshGeom(mesh_path=mesh_path, mesh=mesh, pos=pos, quat=quat, scale=scale)
geoms.append(geom)
# Shift everything so that the reference of the body is at the very center of the composed
# object. This is very important.
off_center_pos, _ = compute_pos_and_size(geoms)
for geom in geoms:
geom.pos -= off_center_pos
# Ensure that the object origin is exactly at the center.
assert np.allclose(compute_pos_and_size(geoms)[0], 0.0)
# Resize object.
if object_size is not None:
_, size = compute_pos_and_size(geoms)
# Apply global scale (so that ratio is not changed and longest side is exactly
# object_size).
ratio = object_size / np.max(size)
for geom in geoms:
geom.scale *= ratio
geom.pos *= ratio
geoms_str = "\n".join([g.to_geom_xml(name, idx) for idx, g in enumerate(geoms)])
meshes_str = "\n".join([g.to_mesh_xml(name, idx) for idx, g in enumerate(geoms)])
xml_source = f"""
<mujoco>
<asset>
{meshes_str}
</asset>
<worldbody>
<body name="{name}" pos="0 0 0">
{geoms_str}
<joint name="{name}:joint" type="free"/>
</body>
</worldbody>
</mujoco>
"""
return MujocoXML.from_string(xml_source)
def make_world_xml(cls, *, contact_params: dict, mujoco_timestep: float,
**kwargs):
return (MujocoXML.parse(cls.BASE_XML).set_objects_attr(
tag="option", timestep=mujoco_timestep).set_objects_attr(
tag="size", **contact_params).add_default_compiler_directive())
def make_openai_block(name: str, object_size: np.ndarray) -> MujocoXML:
""" Creates a block with OPENAI letters on it faces.
:param name: The name of the block
:param object_size: The size of the block (3-dimensional). This is half-size as per Mujoco
convention
"""
default_object_size = 0.0285
default_letter_offset = 0.0009
# scale face meshes properly
scale = object_size / default_object_size
letter_offset = default_letter_offset * scale
def to_str(x: np.ndarray):
return " ".join(map(str, x.tolist()))
face_pos = {
"top": {
"body": to_str(np.array([0, 0, object_size[2]])),
"geom": to_str(np.array([0, 0, -letter_offset[2]])),
},
"bottom": {
"body": to_str(np.array([0, 0, -object_size[2]])),
"geom": to_str(np.array([0, 0, letter_offset[2]])),
},
"back": {
"body": to_str(np.array([0, object_size[1], 0])),
"geom": to_str(np.array([0, -letter_offset[1], 0])),
},
"right": {
"body": to_str(np.array([object_size[0], 0, 0])),
"geom": to_str(np.array([-letter_offset[0], 0, 0])),
},
"front": {
"body": to_str(np.array([0, -object_size[1], 0])),
"geom": to_str(np.array([0, letter_offset[1], 0])),
},
"left": {
"body": to_str(np.array([-object_size[0], 0, 0])),
"geom": to_str(np.array([letter_offset[0], 0, 0])),
},
}
face_euler = {
"top": to_str(np.array([np.pi / 2, 0, np.pi / 2])),
"bottom": to_str(np.array([np.pi / 2, 0, np.pi / 2])),
"back": to_str(np.array([0, 0, np.pi / 2])),
"right": to_str(np.array([0, 0, 0])),
"front": to_str(np.array([0, 0, -np.pi / 2])),
"left": to_str(np.array([0, 0, np.pi])),
}
def face_xml(_name: str, _face: str, _c: str):
xml = f"""
<body name="{_face}:{_name}" pos="{face_pos[_face]['body']}">
<geom name="letter_{_c}:{_name}" mesh="{_name}{_c}" euler="{face_euler[_face]}"
pos="{face_pos[_face]['geom']}" type="mesh" material="{_name}letter"
conaffinity="0" contype="0" />
</body>
"""
return xml
size_string = " ".join(map(str, list(object_size)))
scale_string = " ".join(map(str, list(scale)))
xml_source = f"""
<mujoco>
<asset>
<material name="{name}letter" specular="1" shininess="0.3" rgba="1 1 1 1"/>
<mesh name="{name}O" file="{ASSETS_DIR}/stls/openai_cube/O.stl"
scale="{scale_string}" />
<mesh name="{name}P" file="{ASSETS_DIR}/stls/openai_cube/P.stl"
scale="{scale_string}" />
<mesh name="{name}E" file="{ASSETS_DIR}/stls/openai_cube/E.stl"
scale="{scale_string}" />
<mesh name="{name}N" file="{ASSETS_DIR}/stls/openai_cube/N.stl"
scale="{scale_string}" />
<mesh name="{name}A" file="{ASSETS_DIR}/stls/openai_cube/A.stl"
scale="{scale_string}" />
<mesh name="{name}I" file="{ASSETS_DIR}/stls/openai_cube/I.stl"
scale="{scale_string}" />
</asset>
<worldbody>
<body name="{name}">
<geom name="{name}" size="{size_string}" type="box"
rgba="0.0 0.0 0.0 0.0" material="block_mat"/>
<joint name="{name}:joint" type="free"/>
{face_xml(name, "top", "O")}
{face_xml(name, "bottom", "P")}
{face_xml(name, "back", "E")}
{face_xml(name, "right", "N")}
{face_xml(name, "front", "A")}
{face_xml(name, "left", "I")}
</body>
</worldbody>
</mujoco>
"""
return MujocoXML.from_string(xml_source)