key_handler = KeyboardHandler(env=env, cam_body_id=cam_body_id)
env.viewer.add_keypress_callback("any", key_handler.on_press)
env.viewer.add_keyup_callback("any", key_handler.on_release)
env.viewer.add_keyrepeat_callback("any", key_handler.on_press)
# just spin to let user interact with glfw window
spin_count = 0
while True:
action = np.zeros(env.action_dim)
obs, reward, done, _ = env.step(action)
env.render()
spin_count += 1
if spin_count % 500 == 0:
# convert from world coordinates to file coordinates (xml subtree)
camera_pos = np.array(env.sim.model.body_pos[cam_body_id])
camera_quat = T.convert_quat(env.sim.model.body_quat[cam_body_id], to='xyzw')
world_camera_pose = T.make_pose(camera_pos, T.quat2mat(camera_quat))
file_camera_pose = world_in_file.dot(world_camera_pose)
# TODO: Figure out why numba causes black screen of death (specifically, during mat2pose --> mat2quat call below)
camera_pos, camera_quat = T.mat2pose(file_camera_pose)
camera_quat = T.convert_quat(camera_quat, to='wxyz')
print("\n\ncurrent camera tag you should copy")
cam_tree.set("pos", "{} {} {}".format(camera_pos[0], camera_pos[1], camera_pos[2]))
cam_tree.set("quat", "{} {} {} {}".format(camera_quat[0], camera_quat[1], camera_quat[2], camera_quat[3]))
print(ET.tostring(cam_tree, encoding="utf8").decode("utf8"))
def _get_observation(self):
"""
Returns an OrderedDict containing observations [(name_string, np.array), ...].
Important keys:
robot-state: contains robot-centric information.
object-state: requires @self.use_object_obs to be True.
contains object-centric information.
image: requires @self.use_camera_obs to be True.
contains a rendered frame from the simulation.
depth: requires @self.use_camera_obs and @self.camera_depth to be True.
contains a rendered depth map from the simulation
"""
di = super()._get_observation()
if self.use_camera_obs:
camera_obs = self.sim.render(
camera_name=self.camera_name,
width=self.camera_width,
height=self.camera_height,
depth=self.camera_depth,
)
if self.camera_depth:
di["image"], di["depth"] = camera_obs
else:
di["image"] = camera_obs
# low-level object information
if self.use_object_obs:
# remember the keys to collect into object info
object_state_keys = []
# for conversion to relative gripper frame
gripper_pose = T.pose2mat((di["eef_pos"], di["eef_quat"]))
world_pose_in_gripper = T.pose_inv(gripper_pose)
for i in range(len(self.item_names_org)):
if self.single_object_mode == 2 and self.object_id != i:
# Skip adding to observations
continue
obj_str = str(self.item_names_org[i]) + "0"
obj_pos = np.array(
self.sim.data.body_xpos[self.obj_body_id[obj_str]])
obj_quat = T.convert_quat(
self.sim.data.body_xquat[self.obj_body_id[obj_str]],
to="xyzw")
di["{}_pos".format(obj_str)] = obj_pos
di["{}_quat".format(obj_str)] = obj_quat
# get relative pose of object in gripper frame
object_pose = T.pose2mat((obj_pos, obj_quat))
rel_pose = T.pose_in_A_to_pose_in_B(object_pose,
world_pose_in_gripper)
rel_pos, rel_quat = T.mat2pose(rel_pose)
di["{}_to_eef_pos".format(obj_str)] = rel_pos
di["{}_to_eef_quat".format(obj_str)] = rel_quat
object_state_keys.append("{}_pos".format(obj_str))
object_state_keys.append("{}_quat".format(obj_str))
object_state_keys.append("{}_to_eef_pos".format(obj_str))
object_state_keys.append("{}_to_eef_quat".format(obj_str))
if self.single_object_mode == 1:
# Zero out other objects observations
for obj_str, obj_mjcf in self.mujoco_objects.items():
if obj_str == self.obj_to_use:
continue
else:
di["{}_pos".format(obj_str)] *= 0.0
di["{}_quat".format(obj_str)] *= 0.0
di["{}_to_eef_pos".format(obj_str)] *= 0.0
di["{}_to_eef_quat".format(obj_str)] *= 0.0
di["object-state"] = np.concatenate(
[di[k] for k in object_state_keys])
return di
def do_physical_simulation(_sim,
_object_list,
_action,
_next_object_list,
_viewer=None,
test_horizon=20000,
reset=True):
state = _sim.get_state()
if reset:
for _obj in _object_list:
if 'custom_table' not in _obj.name:
pos_arr, quat_arr = T.mat2pose(_obj.pose)
obj_qpos_addr = _sim.model.get_joint_qpos_addr(_obj.name)
state.qpos[obj_qpos_addr[0]:obj_qpos_addr[0] + 3] = pos_arr
state.qpos[obj_qpos_addr[0] + 3:obj_qpos_addr[0] +
7] = quat_arr[[3, 0, 1, 2]]
if action['type'] == "pick":
pass
elif action['type'] == "place":
held_obj_idx = get_held_object(_object_list)
_obj = _next_object_list[held_obj_idx]
pos_arr, quat_arr = T.mat2pose(deepcopy(_obj.pose))
pos_arr[2] += 0.02
obj_qpos_addr = _sim.model.get_joint_qpos_addr(_obj.name)
state.qpos[obj_qpos_addr[0]:obj_qpos_addr[0] + 3] = pos_arr
state.qpos[obj_qpos_addr[0] + 3:obj_qpos_addr[0] +
7] = quat_arr[[3, 0, 1, 2]]
_sim.set_state(state)
for _ in range(test_horizon):
_sim.step()
if _viewer is not None:
_viewer.render()
state = _sim.get_state()
mis_posed_object_list = []
sim_object_list = []
for _obj in _next_object_list:
sim_obj = deepcopy(_obj)
if 'custom_table' not in _obj.name:
pos_arr, quat_arr = T.mat2pose(_obj.pose)
obj_qpos_addr = _sim.model.get_joint_qpos_addr(_obj.name)
sim_pos_arr = state.qpos[obj_qpos_addr[0]:obj_qpos_addr[0] + 3]
sim_quat_arr = state.qpos[obj_qpos_addr[0] + 3:obj_qpos_addr[0] +
7]
sim_obj.pose = T.pose2mat(
[sim_pos_arr, sim_quat_arr[[1, 2, 3, 0]]])
dist_diff = np.sqrt(np.sum((pos_arr - sim_pos_arr)**2))
ang_diff = np.arccos(
np.maximum(
np.minimum(
2. * np.sum(quat_arr * sim_quat_arr[[1, 2, 3, 0]])**2 -
1., 1.), -1.))
if dist_diff > 0.05 or ang_diff > np.pi / 20.:
mis_posed_object_list.append(_obj)
sim_object_list.append(sim_obj)
return mis_posed_object_list, sim_object_list
color=obj.color)
if goal_obj is not None:
ax = visualize_trimesh(meshes[goal_obj.mesh_idx],
goal_obj.pose,
ax=ax,
color=goal_obj.color)
plt.show()
arena_model = MujocoXML(
xml_path_completion("arenas/empty_arena.xml"))
for obj in initial_object_list:
if 'custom_table' in obj.name:
table_model = MujocoXML(
xml_path_completion("objects/custom_table.xml"))
table_pos_arr, table_quat_arr = T.mat2pose(obj.pose)
table_body = table_model.worldbody.find(
"./body[@name='custom_table']")
table_body.set("pos", array_to_string(table_pos_arr))
table_body.set(
"quat", array_to_string(table_quat_arr[[3, 0, 1, 2]]))
arena_model.merge(table_model)
else:
if 'arch_box' in obj.name:
obj_xml_path = "objects/arch_box.xml"
elif 'rect_box' in obj.name:
obj_xml_path = "objects/rect_box.xml"
elif 'square_box' in obj.name:
obj_xml_path = "objects/square_box.xml"
elif 'half_cylinder_box' in obj.name:
obj_xml_path = "objects/half_cylinder_box.xml"
def _get_observation(self):
"""
Returns an OrderedDict containing observations [(name_string, np.array), ...].
Important keys:
`'robot-state'`: contains robot-centric information.
`'object-state'`: requires @self.use_object_obs to be True. Contains object-centric information.
`'image'`: requires @self.use_camera_obs to be True. Contains a rendered frame from the simulation.
`'depth'`: requires @self.use_camera_obs and @self.camera_depth to be True.
Contains a rendered depth map from the simulation
Returns:
OrderedDict: Observations from the environment
"""
di = super()._get_observation()
# low-level object information
if self.use_object_obs:
# Get robot prefix
pr = self.robots[0].robot_model.naming_prefix
# remember the keys to collect into object info
object_state_keys = []
# for conversion to relative gripper frame
gripper_pose = T.pose2mat(
(di[pr + "eef_pos"], di[pr + "eef_quat"]))
world_pose_in_gripper = T.pose_inv(gripper_pose)
for i in range(len(self.item_names_org)):
if self.single_object_mode == 2 and self.nut_id != i:
# skip observations
continue
obj_str = str(self.item_names_org[i]) + "0"
obj_pos = np.array(
self.sim.data.body_xpos[self.obj_body_id[obj_str]])
obj_quat = T.convert_quat(
self.sim.data.body_xquat[self.obj_body_id[obj_str]],
to="xyzw")
di["{}_pos".format(obj_str)] = obj_pos
di["{}_quat".format(obj_str)] = obj_quat
object_pose = T.pose2mat((obj_pos, obj_quat))
rel_pose = T.pose_in_A_to_pose_in_B(object_pose,
world_pose_in_gripper)
rel_pos, rel_quat = T.mat2pose(rel_pose)
di["{}_to_{}eef_pos".format(obj_str, pr)] = rel_pos
di["{}_to_{}eef_quat".format(obj_str, pr)] = rel_quat
object_state_keys.append("{}_pos".format(obj_str))
object_state_keys.append("{}_quat".format(obj_str))
object_state_keys.append("{}_to_{}eef_pos".format(obj_str, pr))
object_state_keys.append("{}_to_{}eef_quat".format(
obj_str, pr))
if self.single_object_mode == 1:
# zero out other objs
for obj_str, obj_mjcf in self.mujoco_objects.items():
if obj_str == self.obj_to_use:
continue
else:
di["{}_pos".format(obj_str)] *= 0.0
di["{}_quat".format(obj_str)] *= 0.0
di["{}_to_{}eef_pos".format(obj_str, pr)] *= 0.0
di["{}_to_{}eef_quat".format(obj_str, pr)] *= 0.0
di["object-state"] = np.concatenate(
[di[k] for k in object_state_keys])
return di
mcts = Tree(initial_object_list, np.sum(n_obj_per_mesh_types) * 2, coll_mngr, meshes, contact_points,
contact_faces, rotation_types, _goal_obj=goal_obj, _physcial_constraint_checker=None)
for _ in range(10):
mcts.exploration(0)
print("Planning")
optimized_path = mcts.get_best_path(0)
final_object_list = mcts.Tree.nodes[optimized_path[-1]]['state']
print("Do dynamic simulation")
if len(configuration_list) == 0:
arena_model = MujocoXML(xml_path_completion("arenas/empty_arena.xml"))
for obj in initial_object_list:
if 'custom_table' in obj.name:
table_model = MujocoXML(xml_path_completion("objects/custom_table.xml"))
table_pos_arr, table_quat_arr = T.mat2pose(obj.pose)
table_body = table_model.worldbody.find("./body[@name='custom_table']")
table_body.set("pos", array_to_string(table_pos_arr))
table_body.set("quat", array_to_string(table_quat_arr[[3, 0, 1, 2]]))
arena_model.merge(table_model)
else:
if 'arch_box' in obj.name:
obj_xml_path = "objects/arch_box.xml"
elif 'rect_box' in obj.name:
obj_xml_path = "objects/rect_box.xml"
elif 'square_box' in obj.name:
obj_xml_path = "objects/square_box.xml"
elif 'half_cylinder_box' in obj.name:
obj_xml_path = "objects/half_cylinder_box.xml"
elif 'triangle_box' in obj.name:
obj_xml_path = "objects/triangle_box.xml"