def add_physx_scene(self, scene, render_shapes_with_one_of_flags=(ShapeFlag.VISUALIZATION,), offset=None):
actors = scene.get_dynamic_rigid_actors() + scene.get_static_rigid_actors()
start_index = self.get_start_index_for_next_scene()
for i, actor in enumerate(actors, start=start_index):
for j, shape in enumerate(actor.get_atached_shapes()):
if not self.has_shape_any_of_flags(shape, render_shapes_with_one_of_flags):
continue
if shape.get_geometry_type() == GeometryType.PLANE: # plane is ignored as there is a grid in meshcat
continue
self.vis_shape(i, j).set_object(self._get_shape_geometry(shape), self._get_shape_material(shape))
self.vis_shape(i, j).set_transform(pose_to_transformation_matrix(shape.get_local_pose()))
if self.show_frames:
self.vis_frame(i).set_object(g.triad(self.frame_scale))
self.actors_and_offsets.append((actors, offset, start_index))
def runTest(self):
self.vis.delete()
v = self.vis["shapes"]
v.set_transform(tf.translation_matrix([1., 0, 0]))
v["box"].set_object(g.Box([1.0, 0.2, 0.3]))
v["box"].delete()
v["box"].set_object(g.Box([0.1, 0.2, 0.3]))
v["box"].set_transform(tf.translation_matrix([0.05, 0.1, 0.15]))
v["cylinder"].set_object(g.Cylinder(0.2, 0.1), g.MeshLambertMaterial(color=0x22dd22))
v["cylinder"].set_transform(tf.translation_matrix([0, 0.5, 0.1]).dot(tf.rotation_matrix(-np.pi / 2, [1, 0, 0])))
v["sphere"].set_object(g.Mesh(g.Sphere(0.15), g.MeshLambertMaterial(color=0xff11dd)))
v["sphere"].set_transform(tf.translation_matrix([0, 1, 0.15]))
v["ellipsoid"].set_object(g.Ellipsoid([0.3, 0.1, 0.1]))
v["ellipsoid"].set_transform(tf.translation_matrix([0, 1.5, 0.1]))
v["transparent_ellipsoid"].set_object(g.Mesh(
g.Ellipsoid([0.3, 0.1, 0.1]),
g.MeshLambertMaterial(color=0xffffff,
opacity=0.5)))
v["transparent_ellipsoid"].set_transform(tf.translation_matrix([0, 2.0, 0.1]))
v = self.vis["meshes/valkyrie/head"]
v.set_object(g.Mesh(
g.ObjMeshGeometry.from_file(os.path.join(meshcat.viewer_assets_path(), "data/head_multisense.obj")),
g.MeshLambertMaterial(
map=g.ImageTexture(
image=g.PngImage.from_file(os.path.join(meshcat.viewer_assets_path(), "data/HeadTextureMultisense.png"))
)
)
))
v.set_transform(tf.translation_matrix([0, 0.5, 0.5]))
v = self.vis["meshes/convex"]
v["obj"].set_object(g.Mesh(g.ObjMeshGeometry.from_file(os.path.join(meshcat.viewer_assets_path(), "../tests/data/mesh_0_convex_piece_0.obj"))))
v["stl_ascii"].set_object(g.Mesh(g.StlMeshGeometry.from_file(os.path.join(meshcat.viewer_assets_path(), "../tests/data/mesh_0_convex_piece_0.stl_ascii"))))
v["stl_ascii"].set_transform(tf.translation_matrix([0, -0.5, 0]))
v["stl_binary"].set_object(g.Mesh(g.StlMeshGeometry.from_file(os.path.join(meshcat.viewer_assets_path(), "../tests/data/mesh_0_convex_piece_0.stl_binary"))))
v["stl_binary"].set_transform(tf.translation_matrix([0, -1, 0]))
v["dae"].set_object(g.Mesh(g.DaeMeshGeometry.from_file(os.path.join(meshcat.viewer_assets_path(), "../tests/data/mesh_0_convex_piece_0.dae"))))
v["dae"].set_transform(tf.translation_matrix([0, -1.5, 0]))
v = self.vis["points"]
v.set_transform(tf.translation_matrix([0, 2, 0]))
verts = np.random.rand(3, 1000000)
colors = verts
v["random"].set_object(g.PointCloud(verts, colors))
v["random"].set_transform(tf.translation_matrix([-0.5, -0.5, 0]))
v = self.vis["lines"]
v.set_transform(tf.translation_matrix(([-2, -3, 0])))
vertices = np.random.random((3, 10)).astype(np.float32)
v["line_segments"].set_object(g.LineSegments(g.PointsGeometry(vertices)))
v["line"].set_object(g.Line(g.PointsGeometry(vertices)))
v["line"].set_transform(tf.translation_matrix([0, 1, 0]))
v["line_loop"].set_object(g.LineLoop(g.PointsGeometry(vertices)))
v["line_loop"].set_transform(tf.translation_matrix([0, 2, 0]))
v["line_loop_with_material"].set_object(g.LineLoop(g.PointsGeometry(vertices), g.LineBasicMaterial(color=0xff0000)))
v["line_loop_with_material"].set_transform(tf.translation_matrix([0, 3, 0]))
colors = vertices # Color each line by treating its xyz coordinates as RGB colors
v["line_with_vertex_colors"].set_object(g.Line(g.PointsGeometry(vertices, colors), g.LineBasicMaterial(vertexColors=True)))
v["line_with_vertex_colors"].set_transform(tf.translation_matrix([0, 4, 0]))
v["triad"].set_object(g.LineSegments(
g.PointsGeometry(position=np.array([
[0, 0, 0], [1, 0, 0],
[0, 0, 0], [0, 1, 0],
[0, 0, 0], [0, 0, 1]]).astype(np.float32).T,
color=np.array([
[1, 0, 0], [1, 0.6, 0],
[0, 1, 0], [0.6, 1, 0],
[0, 0, 1], [0, 0.6, 1]]).astype(np.float32).T
),
g.LineBasicMaterial(vertexColors=True)))
v["triad"].set_transform(tf.translation_matrix(([0, 5, 0])))
v["triad_function"].set_object(g.triad(0.5))
v["triad_function"].set_transform(tf.translation_matrix([0, 6, 0]))
def draw_tree(tree, vis, prefix="", draw_regions=False):
# Given a scene tree (nx.DiGraph), draw it in the
# specified meshcat visualizer.
# Draw the scene geometry flat, to keep TFs easy.
name_prefix = prefix + "scene"
vis[name_prefix].delete()
k = 0
for node in tree.nodes:
name = name_prefix + "/%s_%03d" % (node.__class__.__name__, k)
if node.geometry is not None:
color = node.geometry_color
alpha = 1.0
vis[name].set_object(
node.geometry,
meshcat_geom.MeshLambertMaterial(color=color, opacity=alpha, transparent=(alpha != 1.))
)
tf = node.tf.GetAsMatrix4()
geom_tf = node.geometry_tf.GetAsMatrix4()
tf = tf.dot(geom_tf)
tf[:3, :3] = tf[:3, :3].dot(np.diag(node.geometry_scale))
print(tf)
vis[name].set_transform(tf)
k += 1
# Draw the tree structure.
tree_prefix = prefix + "tree"
vis[tree_prefix].delete()
k = 0
for node in tree.nodes:
name = tree_prefix + "/" + node.__class__.__name__ + "_%03d" % k
k += 1
# Draw node as randomly colored sphere
color = random.randint(0, 0xFFFFFF)
alpha = 0.5
vis[name]["triad"].set_object(
meshcat_geom.triad(scale=0.1)
)
vis[name]["sphere"].set_object(
meshcat_geom.Sphere(0.01),
meshcat_geom.MeshToonMaterial(color=color, opacity=alpha, transparent=(alpha != 1.))
)
vis[name].set_transform(node.tf.GetAsMatrix4())
# Draw children
verts = []
for child in tree.successors(node):
# Draw link to child
verts.append(node.tf.translation()),
verts.append(child.tf.translation())
if len(verts) > 0:
verts = np.vstack(verts).T
# Don't want this as a direct child or it'll inherit the transform
vis[name + "_child_connections"].set_object(
meshcat_geom.Line(meshcat_geom.PointsGeometry(verts),
meshcat_geom.LineBasicMaterial(linewidth=50, color=color)))
if draw_regions:
# Draw the child regions for each child
if isinstance(node, (AndNode, OrNode, RepeatingSetNode)):
for info_k, child_info in enumerate(node.child_infos):
region_name = "child_region_%03d" % info_k
lb = child_info.child_xyz_bounds.xyz_min
ub = child_info.child_xyz_bounds.xyz_max
vis[name][region_name].set_object(
meshcat_geom.Box(ub - lb),
meshcat_geom.MeshToonMaterial(color=0x111111, opacity=0.1, transparent=True)
)
tf = RigidTransform(p=(ub+lb)/2)
vis[name][region_name].set_transform(tf.GetAsMatrix4())
def main():
# load dynamics model
model_dict = load_autoencoder_model()
model = model_dict['model_dy']
model_ae = model_dict['model_ae']
visual_observation_function = model_dict['visual_observation_function']
config = model.config
env_config = load_yaml(
os.path.join(get_project_root(),
'experiments/exp_18_box_on_side/config.yaml'))
env_config['env']['observation']['depth_int16'] = True
n_history = config['train']['n_history']
# create the environment
# create the environment
env = DrakePusherSliderEnv(env_config)
env.reset()
# create another environment for doing rollouts
env2 = DrakePusherSliderEnv(env_config, visualize=False)
env2.reset()
action_function = ActionFunctionFactory.function_from_config(config)
observation_function = ObservationFunctionFactory.drake_pusher_position_3D(
config)
episode = OnlineEpisodeReader()
mpc_input_builder = DynamicsModelInputBuilder(
observation_function=observation_function,
visual_observation_function=visual_observation_function,
action_function=action_function,
episode=episode)
vis = meshcat_utils.make_default_visualizer_object()
vis.delete()
initial_cond = get_initial_state()
reset_environment(env, initial_cond['q_pusher'], initial_cond['q_slider'])
obs_init = env.get_observation()
# visualize starting position of the object
print("obs_init.keys()", obs_init.keys())
print("obs_init['slider']['position']", obs_init['slider']['position'])
T = DrakePusherSliderEnv.object_position_from_observation(obs_init)
vis['start_pose'].set_object(triad(scale=0.1))
vis['state_pose'].set_transform(T)
#### ROLLOUT USING LEARNED MODEL + GROUND TRUTH ACTIONS ############
reset_environment(env, initial_cond['q_pusher'], initial_cond['q_slider'])
# add just some large number of these
episode.clear()
for i in range(n_history):
action_zero = np.zeros(2)
obs_tmp = env.get_observation()
episode.add_observation_action(obs_tmp, action_zero)
#### ROLLOUT THE ACTION SEQUENCE USING THE SIMULATOR ##########
# rollout single action sequence using the simulator
gt_rollout_data = env_utils.rollout_action_sequence(
env, initial_cond['action_sequence'].cpu().numpy())
env_obs_rollout_gt = gt_rollout_data['observations']
gt_rollout_episode = gt_rollout_data['episode_reader']
for i, env_obs in enumerate(gt_rollout_data['observations']):
T = DrakePusherSliderEnv.object_position_from_observation(env_obs)
vis_name = "GT_trajectory/%d" % (i)
vis[vis_name].set_object(triad(scale=0.1))
vis[vis_name].set_transform(T)
action_state_gt = mpc_input_builder.get_action_state_tensors(
start_idx=0, num_timesteps=N, episode=gt_rollout_episode)
state_rollout_gt = action_state_gt['states']
action_rollout_gt = action_state_gt['actions']
z_object_rollout_gt = model.compute_z_state(
state_rollout_gt)['z_object_flat']
print('state_rollout_gt.shape', state_rollout_gt.shape)
print("z_object_rollout_gt.shape", z_object_rollout_gt.shape)
def goal_func(obs_tmp):
state_tmp = mpc_input_builder.get_state_input_single_timestep(
{'observation': obs_tmp})['state']
return model.compute_z_state(
state_tmp.unsqueeze(0))['z_object'].flatten()
# using the vision model to get "goal" keypoints
z_object_goal = goal_func(env_obs_rollout_gt[-1])
z_object_goal_np = torch_utils.cast_to_numpy(z_object_goal)
# input("press Enter to continue")
#### ROLLOUT USING LEARNED MODEL + GROUND TRUTH ACTIONS ############
reset_environment(env, initial_cond['q_pusher'], initial_cond['q_slider'])
# add just some large number of these
episode.clear()
for i in range(n_history):
action_zero = np.zeros(2)
obs_tmp = env.get_observation()
episode.add_observation_action(obs_tmp, action_zero)
# [n_history, state_dim]
idx = episode.get_latest_idx()
dyna_net_input = mpc_input_builder.get_dynamics_model_input(
idx, n_history=n_history)
state_init = dyna_net_input['states'].cuda() # [n_history, state_dim]
action_init = dyna_net_input['actions'] # [n_history, action_dim]
print("state_init.shape", state_init.shape)
print("action_init.shape", action_init.shape)
print("n_history", n_history)
action_seq_gt_torch = initial_cond['action_sequence']
action_input = torch.cat(
(action_init[:(n_history - 1)], action_seq_gt_torch), dim=0).cuda()
print("action_input.shape", action_input.shape)
# rollout using the ground truth actions and learned model
# need to add the batch dim to do that
z_init = model.compute_z_state(state_init)['z']
rollout_pred = rollout_model(state_init=z_init.unsqueeze(0),
action_seq=action_input.unsqueeze(0),
dynamics_net=model,
compute_debug_data=True)
state_pred_rollout = rollout_pred['state_pred'].squeeze(0)
print("state_pred_rollout.shape", state_pred_rollout.shape)
# input("press Enter to continue")
# check L2 distance between predicted and actual
# basically comparing state_pred_rollout and state_rollout_gt
print("state_rollout_gt[-1]\n", state_rollout_gt[-1])
print("state_pred_rollout[-1]\n", state_pred_rollout[-1])
index_dict = get_object_and_robot_state_indices(config)
object_indices = index_dict['object_indices']
# reset the environment and use the MPC controller to stabilize this
# now setup the MPC to try to stabilize this . . . .
reset_environment(env, initial_cond['q_pusher'], initial_cond['q_slider'])
episode.clear()
# add just some large number of these
for i in range(n_history):
action_zero = np.zeros(2)
obs_tmp = env.get_observation()
episode.add_observation_action(obs_tmp, action_zero)
# input("press Enter to continue")
# make a planner config
planner_config = copy.copy(config)
config_tmp = load_yaml(
os.path.join(get_project_root(),
'experiments/drake_pusher_slider/eval_config.yaml'))
planner_config['mpc'] = config_tmp['mpc']
planner_config['mpc']['mppi']['terminal_cost_only'] = TERMINAL_COST_ONLY
planner_config['mpc']['random_shooting'][
'terminal_cost_only'] = TERMINAL_COST_ONLY
planner = None
if PLANNER_TYPE == "random_shooting":
planner = RandomShootingPlanner(planner_config)
elif PLANNER_TYPE == "mppi":
planner = PlannerMPPI(planner_config)
else:
raise ValueError("unknown planner type: %s" % (PLANNER_TYPE))
mpc_out = None
action_seq_mpc = None
state_pred_mpc = None
counter = -1
while True:
counter += 1
print("\n\n-----Running MPC Optimization: Counter (%d)-------" %
(counter))
obs_cur = env.get_observation()
episode.add_observation_only(obs_cur)
if counter == 0 or REPLAN:
print("replanning")
####### Run the MPC ##########
# [1, state_dim]
n_look_ahead = N - counter
if USE_FIXED_MPC_HORIZON:
n_look_ahead = MPC_HORIZON
elif USE_SHORT_HORIZON_MPC:
n_look_ahead = min(MPC_HORIZON, N - counter)
if n_look_ahead == 0:
break
start_idx = counter
end_idx = counter + n_look_ahead
print("start_idx:", start_idx)
print("end_idx:", end_idx)
print("n_look_ahead", n_look_ahead)
# start_time = time.time()
# idx of current observation
idx = episode.get_latest_idx()
mpc_start_time = time.time()
mpc_input_data = mpc_input_builder.get_dynamics_model_input(
idx, n_history=n_history)
state_cur = mpc_input_data['states']
action_his = mpc_input_data['actions']
if SEED_WITH_NOMINAL_ACTIONS:
action_seq_rollout_init = action_seq_gt_torch[
start_idx:end_idx]
else:
if mpc_out is not None:
action_seq_rollout_init = mpc_out['action_seq'][1:]
print("action_seq_rollout_init.shape",
action_seq_rollout_init.shape)
if action_seq_rollout_init.shape[0] < n_look_ahead:
num_steps = n_look_ahead - action_seq_rollout_init.shape[
0]
action_seq_zero = torch.zeros([num_steps, 2])
action_seq_rollout_init = torch.cat(
(action_seq_rollout_init, action_seq_zero), dim=0)
print("action_seq_rollout_init.shape",
action_seq_rollout_init.shape)
else:
action_seq_rollout_init = None
# run MPPI
z_cur = None
with torch.no_grad():
z_cur = model.compute_z_state(
state_cur.unsqueeze(0).cuda())['z'].squeeze(0)
if action_seq_rollout_init is not None:
n_look_ahead = action_seq_rollout_init.shape[0]
obs_goal = None
print("z_object_rollout_gt.shape", z_object_rollout_gt.shape)
if TRAJECTORY_GOAL:
obs_goal = z_object_rollout_gt[start_idx:end_idx]
print("n_look_ahead", n_look_ahead)
print("obs_goal.shape", obs_goal.shape)
# add the final goal state on as needed
if obs_goal.shape[0] < n_look_ahead:
obs_goal_final = z_object_rollout_gt[-1].unsqueeze(0)
num_repeat = n_look_ahead - obs_goal.shape[0]
obs_goal_final_expand = obs_goal_final.expand(
[num_repeat, -1])
obs_goal = torch.cat((obs_goal, obs_goal_final_expand),
dim=0)
else:
obs_goal = z_object_rollout_gt[-1]
obs_goal = torch_utils.cast_to_numpy(obs_goal)
print("obs_goal.shape", obs_goal.shape)
seed = 1
set_seed(seed)
mpc_out = planner.trajectory_optimization(
state_cur=z_cur,
action_his=action_his,
obs_goal=obs_goal,
model_dy=model,
action_seq_rollout_init=action_seq_rollout_init,
n_look_ahead=n_look_ahead,
eval_indices=object_indices,
rollout_best_action_sequence=True,
verbose=True,
add_grid_action_samples=True,
)
print("MPC step took %.4f seconds" %
(time.time() - mpc_start_time))
action_seq_mpc = torch_utils.cast_to_numpy(mpc_out['action_seq'])
state_pred_mpc = torch_utils.cast_to_numpy(mpc_out['state_pred'])
# Rollout with ground truth simulator dynamics
env2.set_simulator_state_from_observation_dict(
env2.get_mutable_context(), obs_cur)
obs_mpc_gt = env_utils.rollout_action_sequence(
env2, action_seq_mpc)['observations']
vis['mpc_3D'].delete()
vis['mpc_GT_3D'].delete()
L = len(obs_mpc_gt)
print("L", L)
if L == 0:
break
for i in range(L):
# ground truth rollout of the MPC action_seq
name = "mpc_GT_3D/%d" % (i)
T_W_obj = DrakePusherSliderEnv.object_position_from_observation(
obs_mpc_gt[i])
vis[name].set_object(triad(scale=0.1))
vis[name].set_transform(T_W_obj)
action_cur = action_seq_mpc[0]
print("action_cur", action_cur)
print("action_GT", initial_cond['action'])
input("press Enter to continue")
# add observation actions to the episode
obs_cur = env.get_observation()
episode.replace_observation_action(obs_cur, action_cur)
# step the simulator
env.step(action_cur)
# update the trajectories, in case we aren't replanning
action_seq_mpc = action_seq_mpc[1:]
state_pred_mpc = state_pred_mpc[1:]
pose_error = compute_pose_error(env_obs_rollout_gt[-1], obs_cur)
print("position_error: %.3f" % (pose_error['position_error']))
print("angle error degrees: %.3f" %
(pose_error['angle_error_degrees']))
obs_final = env.get_observation()
pose_error = compute_pose_error(env_obs_rollout_gt[-1], obs_final)
print("position_error: %.3f" % (pose_error['position_error']))
print("angle error degrees: %.3f" % (pose_error['angle_error_degrees']))
def draw_scene_tree_structure_meshcat(scene_tree,
prefix="scene_tree",
zmq_url=None,
alpha=0.775,
node_sphere_size=0.05,
linewidth=2,
with_triad=True,
quiet=True,
color_by_score=None,
delete=True):
# Color by score can be a tuple of min, max score. It'll go from red at min score
# to blue at max score.
# Do actual drawing in meshcat.
if quiet:
with open(os.devnull, 'w') as devnull:
with contextlib.redirect_stdout(devnull):
vis = meshcat.Visualizer(
zmq_url=zmq_url or "tcp://127.0.0.1:6000")
else:
vis = meshcat.Visualizer(zmq_url=zmq_url or "tcp://127.0.0.1:6000")
if delete:
vis[prefix].delete()
# Assign functionally random colors to each new node
# type we discover, or color my their scores.
node_class_to_color_dict = {}
cmap = plt.cm.get_cmap('jet')
cmap_counter = 0.
k = 0
for node in scene_tree.nodes:
children, rules = scene_tree.get_children_and_rules(node)
#
if color_by_score is not None:
assert len(color_by_score
) == 2, "Color by score should be a tuple of (min, max)"
score = node.score_child_set(children)
print("Node score: ", score)
score = (score - color_by_score[0]) / (color_by_score[1] -
color_by_score[0])
score = 1. - np.clip(score.item(), 0., 1.)
color = rgb_2_hex(cmap(score))
#color = 0x555555
else:
# Draw this node
node_type_string = node.__class__.__name__
if node_type_string in node_class_to_color_dict.keys():
color = node_class_to_color_dict[node_type_string]
else:
color = rgb_2_hex(cmap(cmap_counter))
node_class_to_color_dict[node_type_string] = color
cmap_counter = np.fmod(cmap_counter + np.pi * 2., 1.)
vis[prefix][node.name + "%d/sphere" % k].set_object(
meshcat_geom.Sphere(node_sphere_size),
meshcat_geom.MeshToonMaterial(color=color,
opacity=alpha,
transparent=(alpha != 1.),
depthTest=False))
if with_triad:
vis[prefix][node.name + "%d/triad" % k].set_object(
meshcat_geom.triad(scale=node_sphere_size * 5.))
tf = node.tf.cpu().detach().numpy()
vis[prefix][node.name + "%d" % k].set_transform(tf)
# Draw connections to each child
for child, rule in zip(children, rules):
verts = []
verts.append(node.tf[:3, 3].cpu().detach().numpy())
verts.append(child.tf[:3, 3].cpu().detach().numpy())
verts = np.vstack(verts).T
if color_by_score is not None:
score = rule.score_child(node, child)
print("Rule score: ", score)
score = (score - color_by_score[0]) / (color_by_score[1] -
color_by_score[0])
score = 1. - np.clip(score.item(), 0., 1.)
color = rgb_2_hex(cmap(score))
vis[prefix][node.name + "_to_" + child.name].set_object(
meshcat_geom.Line(
meshcat_geom.PointsGeometry(verts),
meshcat_geom.LineBasicMaterial(linewidth=linewidth,
color=color,
depthTest=False)))
k += 1