def get_dataset_paths(dataset_name):
if dataset_name == "push_box_hardware":
episodes_root = os.path.join(get_data_ssd_root(),
"dataset/push_box_hardware")
episodes_config = load_yaml(
os.path.join(
get_project_root(),
'experiments/exp_22_push_box_hardware/push_box_hardware_episodes_config.yaml'
))
transporter_model_chkpt = None
dense_descriptor_model_chkpt = "/home/manuelli/data/key_dynam/dev/experiments/22/dataset_push_box_string_pull/trained_models/perception/dense_descriptors/data_aug_2020-07-02-02-39-27-400442/net_best_model.pth"
return {
'dataset_name': dataset_name,
'dataset_root': episodes_root,
'episodes_config': episodes_config,
'main_camera_name': 'd415_01',
'dense_descriptor_camera_list': ['d415_01', 'd415_02'],
'transporter_model_chkpt': transporter_model_chkpt,
'dense_descriptor_model_chkpt': dense_descriptor_model_chkpt,
}
elif dataset_name == "push_box_string_pull":
episodes_root = os.path.join(get_data_ssd_root(),
"dataset/push_box_string_pull")
episodes_config = load_yaml(
os.path.join(
get_project_root(),
'experiments/exp_22_push_box_hardware/push_box_string_pull_episodes_config.yaml'
))
transporter_model_chkpt = None
dense_descriptor_model_chkpt = "/home/manuelli/data/key_dynam/dev/experiments/22/dataset_push_box_string_pull/trained_models/perception/dense_descriptors/data_aug_2020-07-02-02-39-27-400442/net_best_model.pth"
return {
'dataset_name': dataset_name,
'dataset_root': episodes_root,
'episodes_config': episodes_config,
'main_camera_name': 'd415_01',
'dense_descriptor_camera_list': ['d415_01', 'd415_02'],
'transporter_model_chkpt': transporter_model_chkpt,
'dense_descriptor_model_chkpt': dense_descriptor_model_chkpt,
}
else:
raise ValueError("unknown dataset:", dataset_name)
def load_config():
config = load_yaml(
os.path.join(
get_project_root(),
"experiments/exp_22_push_box_hardware/integral_heatmap_3d.yaml"
))
return config
def test_pusher_slider_dataset():
# dataset, config = create_pusher_slider_dataset()
project_root = get_project_root()
config_file = os.path.join(project_root, "experiments/01/config.yaml")
config = load_yaml(config_file)
# new dataset loading approach
episodes = load_episodes_from_config(config)
action_function = ActionFunctionFactory.function_from_config(config)
observation_function = ObservationFunctionFactory.function_from_config(
config)
dataset = MultiEpisodeDataset(config,
action_function=action_function,
observation_function=observation_function,
episodes=episodes,
phase="train")
data = dataset[0] # test the getitem
print("type(data)", type(data))
print("list(data)", list(data))
print(type(data["observations"]))
print("observations.shape", data["observations"].shape)
print("actions.shape", data["actions"].shape)
print("observations", data["observations"])
print("actions", data["actions"])
stats = dataset.compute_dataset_statistics()
print("stats", stats)
def sample_random_mug():
sdf_dir = os.path.join(get_data_root(), "stable/sim_assets/anzu_mugs")
# sdf_file = random.choice(SDFHelper.get_sdf_list(sdf_dir))
mug_list = load_yaml(
os.path.join(get_project_root(), 'experiments/exp_20_mugs/mugs.yaml'))
sdf_file = random.choice(mug_list['corelle_mug-small'])
sdf_file = os.path.join(sdf_dir, sdf_file)
return sdf_file
def load_default_config():
"""
Loads the experiments/05/config.yaml
:return:
:rtype:
"""
config_file = os.path.join(get_project_root(),
'experiments/05/config.yaml')
config = load_yaml(config_file)
return config
def f(q_tmp):
config = load_yaml(
os.path.join(get_project_root(),
'experiments/exp_20_mugs/config.yaml'))
config['dataset']['num_episodes'] = num_episodes_per_thread
out = collect_episodes(config,
output_dir=OUTPUT_DIR,
visualize=False,
debug=False,
run_from_thread=True)
q_tmp.put(out)
def test_pusher_slider_keypoint_dataset():
project_root = get_project_root()
config_file = os.path.join(project_root, "experiments/02/config.yaml")
config = load_yaml(config_file)
config["n_history"] = 1
config["n_roll"] = 0
# new dataset loading approach
episodes = load_episodes_from_config(config)
action_function = ActionFunctionFactory.function_from_config(config)
observation_function = ObservationFunctionFactory.function_from_config(
config)
dataset = MultiEpisodeDataset(config,
action_function=action_function,
observation_function=observation_function,
episodes=episodes,
phase="train")
# dataset, config = create_pusher_slider_keypoint_dataset(config=config)
episode_names = dataset.get_episode_names()
episode_names.sort()
episode_name = episode_names[0]
episode = dataset.episode_dict[episode_name]
obs_raw = episode.get_observation(0)
obs_raw['slider']['angle'] = 0
dataset.observation_function(obs_raw)
print("20 degrees\n\n\n\n")
obs_raw['slider']['angle'] = np.deg2rad(90)
dataset.observation_function(obs_raw)
quit()
data = dataset[0] # test the getitem
print("type(data)", type(data))
print("data.keys()", data.keys())
print(type(data["observations"]))
print("observations.shape", data["observations"].shape)
print("actions.shape", data["actions"].shape)
print("observations", data["observations"])
print("actions", data["actions"])
def load_episodes_from_config(config):
"""
Loads episodes using the path specified in the config
:param config:
:type config:
:return:
:rtype:
"""
data_path = config["dataset"]["data_path"]
if not os.path.isabs(data_path):
data_path = os.path.join(get_project_root(), data_path)
# load the data
print("loading data from disk . . . ")
raw_data = load_pickle(data_path)
print("finished loading data")
episodes = PyMunkEpisodeReader.load_pymunk_episodes_from_raw_data(raw_data)
return episodes
def debug():
save_dir = os.path.join(get_project_root(),
'sandbox/mpc/push_right_box_horizontal')
save_dir = "/home/manuelli/data/key_dynam/sandbox/2020-07-07-20-09-54_push_right_box_horizontal"
plan_msg = load_pickle(os.path.join(save_dir, 'plan_msg.p'),
encoding='latin1')
plan_msg = zmq_utils.convert(plan_msg)
compute_control_action_msg = load_pickle(os.path.join(
save_dir, 'compute_control_action_msg.p'),
encoding='latin1')
compute_control_action_msg = zmq_utils.convert(compute_control_action_msg)
K_matrix = None
T_world_camera = None
if 'K_matrix' in plan_msg:
K_matrix = plan_msg['K_matrix']
if 'T_world_camera' in plan_msg:
T_world_camera = plan_msg['T_world_camera']
d = load_model_and_data(
K_matrix=K_matrix,
T_world_camera=T_world_camera,
)
controller = ZMQController(
config=d['config'],
model_dy=d['model_dy'],
action_function=d['action_function'],
observation_function=d['observation_function'],
visual_observation_function=d['visual_observation_function'],
planner=d['planner'],
debug=True,
zmq_enabled=False,
camera_info=d['camera_info'],
)
controller._on_plan_msg(plan_msg)
controller._on_compute_control_action(compute_control_action_msg)
def construct_dataset_from_config(
config, # dict for global config
phase="train", # str: either "train" or "valid"
episodes=None, # optional dict of episodes to use instead of loading data
):
"""
Construct dataset based on global config
:param config:
:type config:
:return:
:rtype:
"""
assert phase in ["train", "valid"]
data_path = config["dataset"]["data_path"]
if not os.path.isabs(data_path):
data_path = os.path.join(get_project_root(), data_path)
if episodes is None:
# load the data if not passed in
episodes = load_episodes_from_config(config)
action_function = ActionFunctionFactory.function_from_config(config)
observation_function = ObservationFunctionFactory.function_from_config(
config)
dataset = MultiEpisodeDataset(config,
action_function=action_function,
observation_function=observation_function,
episodes=episodes,
phase=phase)
return_data = {
"dataset": dataset,
"action_function": action_function,
"observation_function": observation_function,
"episodes": episodes,
}
return return_data
def main():
start_time = time.time()
config = load_yaml(os.path.join(get_project_root(), 'experiments/drake_pusher_slider/env_config.yaml'))
config['dataset']['num_episodes'] = 1000 # half for train, half for valid
set_seed(500) # just randomly chosen
num_episodes = config['dataset']['num_episodes']
DATASET_NAME = "box_push_%d" %(num_episodes)
OUTPUT_DIR = os.path.join(get_data_ssd_root(), 'dataset', DATASET_NAME)
if not os.path.exists(OUTPUT_DIR):
os.makedirs(OUTPUT_DIR)
collect_episodes(
config,
output_dir=OUTPUT_DIR,
visualize=False,
debug=False)
elapsed = time.time() - start_time
print("Generating and saving dataset to disk took %d seconds" % (int(elapsed)))
def main():
start_time = time.time()
config = load_yaml(
os.path.join(get_project_root(),
'experiments/exp_20_mugs/config.yaml'))
config['dataset']['num_episodes'] = 10
set_seed(500) # just randomly chosen
DATASET_NAME = "mugs_%d" % (config['dataset']['num_episodes'])
OUTPUT_DIR = os.path.join(get_data_root(), 'sandbox', DATASET_NAME)
print("OUTPUT_DIR:", OUTPUT_DIR)
collect_episodes(config,
output_dir=OUTPUT_DIR,
visualize=False,
debug=False)
elapsed = time.time() - start_time
print("Generating and saving dataset to disk took %d seconds" %
(int(elapsed)))
def create_pusher_slider_keypoint_dataset(config=None):
# load some previously generated data
project_root = get_project_root()
if config is None:
config_file = os.path.join(project_root, "experiments/02/config.yaml")
config = load_yaml(config_file)
action_function = ActionFunctionFactory.pusher_velocity
obs_function = ObservationFunctionFactory.pusher_pose_slider_keypoints(
config)
DATA_PATH = os.path.join(
project_root,
"test_data/pusher_slider_10_episodes/2019-10-22-21-30-02-536750.p")
raw_data = load_pickle(DATA_PATH)
episodes = PyMunkEpisodeReader.load_pymunk_episodes_from_raw_data(raw_data)
# create MultiEpisodeDataset
dataset = MultiEpisodeDataset(config,
action_function=action_function,
observation_function=obs_function,
episodes=episodes)
episode = dataset.get_random_episode()
data_0 = episode.get_observation(0)
data_1 = episode.get_observation(1)
print("time 0", data_0["sim_time"])
print("time 1", data_1["sim_time"])
# episode_name = episodes.keys()[0]
# episode = episodes[episode_name]
# data = episode.data
# print("episode.data.keys()", episode.data.keys())
# print("test ", type(data["trajectory"][0].keys()))
# print("test ", data["trajectory"][0].keys())
return dataset, config
def run_interactive_circle_slider():
"""
Launch interactive environment where you can move the pusher around with
the arrow keys
:return:
:rtype:
"""
config_file = os.path.join(get_project_root(),
'experiments/03/config.yaml')
config = load_yaml(config_file)
env = PusherSlider(config=config)
env.reset()
while env._running:
action = env.process_events()
env.step(action)
obs, reward, done, info = env.step(action)
env.render(mode='human')
if True:
print("\n\n\n")
print("slider position", obs['slider']['position'])
print("pusher position", obs['pusher']['position'])
def main():
start_time = time.time()
config = load_yaml(
os.path.join(get_project_root(),
'experiments/exp_18_box_on_side/config.yaml'))
# config['dataset']['num_episodes'] = 500 # half for train, half for valid
config['dataset']['num_episodes'] = 600 # half for train, half for valid
set_seed(500) # just randomly chosen
DATASET_NAME = "dps_box_on_side_%d" % (config['dataset']['num_episodes'])
OUTPUT_DIR = os.path.join(get_data_root(), "dev/experiments/18/data",
DATASET_NAME)
print("OUTPUT_DIR:", OUTPUT_DIR)
collect_episodes(config,
output_dir=OUTPUT_DIR,
visualize=False,
debug=False)
elapsed = time.time() - start_time
print("Generating and saving dataset to disk took %d seconds" %
(int(elapsed)))
def load_model_and_data(
K_matrix=None,
T_world_camera=None,
):
dataset_name = "push_box_hardware"
model_name = "DD_3D/2020-07-02-17-59-21-362337_DD_3D_n_his_2_T_aug"
train_dir = os.path.join(
get_data_root(),
"dev/experiments/22/dataset_push_box_hardware/trained_models/dynamics")
# train_dir = "/home/manuelli/data/key_dynam/dev/experiments/22/dataset_push_box_hardware/trained_models/dynamics"
train_dir = os.path.join(train_dir, model_name)
ckpt_file = os.path.join(train_dir, "net_best_dy_state_dict.pth")
train_config = load_yaml(os.path.join(train_dir, 'config.yaml'))
state_dict = torch.load(ckpt_file)
# build dynamics model
model_dy = build_dynamics_model(train_config)
# print("state_dict.keys()", state_dict.keys())
model_dy.load_state_dict(state_dict)
model_dy = model_dy.eval()
model_dy = model_dy.cuda()
# load the dataset
dataset_paths = get_dataset_paths(dataset_name)
dataset_root = dataset_paths['dataset_root']
episodes_config = dataset_paths['episodes_config']
spatial_descriptor_data = load_pickle(
os.path.join(train_dir, 'spatial_descriptors.p'))
metadata = load_pickle(os.path.join(train_dir, 'metadata.p'))
ref_descriptors = spatial_descriptor_data['spatial_descriptors']
ref_descriptors = torch_utils.cast_to_torch(ref_descriptors).cuda()
# dense descriptor model
model_dd_file = metadata['model_file']
model_dd = torch.load(model_dd_file)
model_dd = model_dd.eval()
model_dd = model_dd.cuda()
camera_name = train_config['dataset']['visual_observation_function'][
'camera_name']
camera_info = None
if (T_world_camera is not None) and (K_matrix is not None):
camera_info = {
"K": K_matrix,
'T_world_camera': T_world_camera,
}
else:
camera_info = get_spartan_camera_info(camera_name)
camera_info['camera_name'] = camera_name
visual_observation_function = \
VisualObservationFunctionFactory.descriptor_keypoints_3D(config=train_config,
camera_name=camera_name,
model_dd=model_dd,
ref_descriptors=ref_descriptors,
K_matrix=camera_info['K'],
T_world_camera=camera_info['T_world_camera'],
)
action_function = ActionFunctionFactory.function_from_config(train_config)
observation_function = ObservationFunctionFactory.function_from_config(
train_config)
#### PLANNER #######
planner = None
# make a planner config
planner_config = copy.copy(train_config)
config_tmp = load_yaml(
os.path.join(get_project_root(),
'experiments/exp_22_push_box_hardware/config_DD_3D.yaml'))
planner_config['mpc'] = config_tmp['mpc']
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))
return {
"model_dy": model_dy,
'config': train_config,
'spatial_descriptor_data': spatial_descriptor_data,
'action_function': action_function,
'observation_function': observation_function,
'visual_observation_function': visual_observation_function,
'planner': planner,
'camera_info': camera_info,
}
def test_gym_API():
"""
Collects two episodes and saves them to disk
Currently no images are being saved
:return:
:rtype:
"""
config_file = os.path.join(get_project_root(),
'experiments/01/config.yaml')
config = load_yaml(config_file)
env = PusherSlider(config=config)
env.setup_environment()
episode_container = None # for namespacing purposes
multi_episode_container = MultiEpisodeContainer()
print("fps", env._fps)
num_episodes = 1
# num_timesteps_per_episode = 40
for episode_idx in range(num_episodes):
env.reset()
obs_prev = env.get_observation()
episode_container = EpisodeContainer()
episode_container.set_config(env.config)
for i in range(env.config['dataset']['num_timesteps']):
env.render(mode='human')
if i == 50:
# save this image out in local directory
img = env.render(mode='rgb_array')
if DEBUG_PRINTS:
print("saving figure to file")
print("type(img)", type(img))
print(img.shape)
print(img.dtype)
pil_img = numpy_to_PIL(img)
pil_img.save('test_PIL.png')
action = np.array([100, 0])
obs, reward, done, info = env.step(action)
episode_container.add_obs_action(obs_prev, action)
obs_prev = obs
if DEBUG_PRINTS:
print("\n\nsim_time = ", env._sim_time)
print("obs sim_time = ", obs['sim_time'])
multi_episode_container.add_episode(episode_container)
if True:
save_file = os.path.join(os.getcwd(), "single_episode_log.p")
episode_container.save_to_file(save_file)
multi_episode_save_file = os.path.join(os.getcwd(),
"multi_episode_log.p")
multi_episode_container.save_to_file(multi_episode_save_file)
def load_config():
return load_yaml(
os.path.join(
get_project_root(),
'experiments/exp_22_push_box_hardware/config_DD_3D.yaml'))
def multiprocess_main(num_episodes=1000, num_threads=4):
set_seed(500) # just randomly chosen
start_time = time.time()
config = load_yaml(
os.path.join(get_project_root(),
'experiments/exp_20_mugs/config.yaml'))
num_episodes_per_thread = math.ceil(num_episodes / num_threads)
num_episodes = num_threads * num_episodes_per_thread
# DATASET_NAME = "mugs_random_colors_%d" % (num_episodes)
# DATASET_NAME = "single_mug_%d"
# DATASET_NAME = "correlle_mug-small_single_color_%d" %(num_episodes)
# DATASET_NAME = "single_corelle_mug_%d" %(num_episodes)
# DATASET_NAME = "correlle_mug-small_many_colors_%d" %(num_episodes)
DATASET_NAME = "correlle_mug-small_many_colors_random_%d" % (num_episodes)
# OUTPUT_DIR = os.path.join(get_data_root(), 'sandbox', DATASET_NAME)
OUTPUT_DIR = os.path.join(get_data_ssd_root(), 'dataset', DATASET_NAME)
print("OUTPUT_DIR:", OUTPUT_DIR)
output_dir = OUTPUT_DIR
if not os.path.exists(output_dir):
os.makedirs(output_dir)
def f(q_tmp):
config = load_yaml(
os.path.join(get_project_root(),
'experiments/exp_20_mugs/config.yaml'))
config['dataset']['num_episodes'] = num_episodes_per_thread
out = collect_episodes(config,
output_dir=OUTPUT_DIR,
visualize=False,
debug=False,
run_from_thread=True)
q_tmp.put(out)
q = Queue()
process_list = []
for i in range(num_threads):
p = Process(target=f, args=(q, ))
p.start()
process_list.append(p)
metadata = {'episodes': {}}
for p in process_list:
while p.is_alive():
p.join(timeout=1)
# empty out the queue
while not q.empty():
out = q.get()
metadata['episodes'].update(out['metadata']['episodes'])
# double check
for p in process_list:
p.join()
time.sleep(1.0)
print("All threads joined")
elapsed = time.time() - start_time
# collect the metadata.yaml files
while not q.empty():
out = q.get()
metadata['episodes'].update(out['metadata']['episodes'])
save_yaml(metadata, os.path.join(OUTPUT_DIR, 'metadata.yaml'))
print("Generating and saving dataset to disk took %d seconds" %
(int(elapsed)))
def load_simple_config():
config_file = os.path.join(get_project_root(), 'config/simple_config.yaml')
config = load_yaml(config_file)
return config
def run_gym_env():
"""
Runs the gym env
:return:
:rtype:
"""
DEBUG_PRINTS = True
USE_PYGAME = True
try:
# setup pygame thing
if USE_PYGAME:
pygame.init()
screen = pygame.display.set_mode(
(640, 480)) # needed for grabbing focus
clock = pygame.time.Clock()
default_action = np.zeros(2)
velocity = 0.2
if USE_PYGAME:
action = process_pygame_events()
else:
action = default_action
config = load_yaml(
os.path.join(get_project_root(),
'experiments/exp_20_mugs/config.yaml'))
env = DrakeMugsEnv(config)
env.reset()
context = env.get_mutable_context()
pos = np.array([0, 0, 0.1])
# quat = transforms3d.euler.euler2quat(np.deg2rad(90), 0, 0)
quat = np.array([1, 0, 0, 0])
q = np.concatenate((quat, pos))
env.set_object_position(context=context, q=q)
# set the box pose
# context = env.get_mutable_context()
# pos = np.array([1.56907481e-04, 1.11390697e-06, 5.11972761e-02])
# quat = np.array([ 7.13518047e-01, -6.69765583e-07, -7.00636851e-01, -6.82079212e-07])
# q_slider = np.concatenate((quat, pos))
# env.set_slider_position(context, q=q_slider)
env.simulator.set_target_realtime_rate(1.0)
# move box araound
# context = env.get_mutable_context()
# q_slider = [-0.05, 0, 0.03]
num_model_instances = env.diagram_wrapper.mbp.num_model_instances()
print("num_model_instances", num_model_instances)
print("num_positions", env.diagram_wrapper.mbp.num_positions())
label_db = env.diagram_wrapper.get_label_db()
print("label db:", label_db.all())
mask_db = env.diagram_wrapper.get_labels_to_mask()
print("mask db:", mask_db.all())
# context = env.get_mutable_context()
#
# # set the position of pusher
# q_pusher = [0.2,0.2]
# mbp = env.diagram_wrapper.mbp
# mbp_context = env.diagram.GetMutableSubsystemContext(mbp, context)
# mbp.SetPositions(mbp_context, env.diagram_wrapper.models['pusher'], q_pusher)
camera_names = list(
config['env']['rgbd_sensors']['sensor_list'].keys())
camera_names.sort()
image_vis = ImageVisualizer(len(camera_names), 1)
print("running sim")
while True:
if USE_PYGAME:
action = velocity * process_pygame_events()
else:
action = default_action
# print("action:", action)
obs, reward, done, info = env.step(action)
# print("obs\n", obs)
# visualize RGB images in matplotlib
for idx, camera_name in enumerate(camera_names):
rgb_image = obs['images'][camera_name]['rgb']
image_vis.draw_image(idx, 0, rgb_image)
image_vis.visualize_interactive()
# # print unique depth values
# depth_32F = obs['images']['camera_0']['depth_32F']
# print("unique depth_32F vals", np.unique(depth_32F))
# # print unique depth values
# depth_16U = obs['images']['camera_0']['depth_16U']
# print("unique depth_16U vals", np.unique(depth_16U))
# build simulator
except KeyboardInterrupt:
pygame.quit()
plt.close()
def test_dynanet_mlp():
# just try doing a single forward pass
dataset, config = create_pusher_slider_dataset()
stats = dataset.compute_dataset_statistics()
n_history = config["train"]["n_history"]
# obs_mean_repeat = stats['observations']['mean'].repeat(n_history, 1)
# obs_std_repeat = stats['observations']['std'].repeat(n_history, 1)
obs_mean_repeat = stats['observations']['mean']
obs_std_repeat = stats['observations']['std']
observations_normalizer = DataNormalizer(obs_mean_repeat, obs_std_repeat)
# action_mean_repeat = stats['actions']['mean'].repeat(n_history, 1)
# action_std_repeat = stats['actions']['std'].repeat(n_history, 1)
action_mean_repeat = stats['actions']['mean']
action_std_repeat = stats['actions']['std']
actions_normalizer = DataNormalizer(action_mean_repeat, action_std_repeat)
config["dataset"]["state_dim"] = 5
config["dataset"]["action_dim"] = 2
model = DynaNetMLP(config)
# print summary of model before adding new modules
print("\n\n -----summary of model BEFORE adding normalization modules")
print("num trainable parameters", count_trainable_parameters(model))
print("num non-trainable parameters ",
count_non_trainable_parameters(model))
print("\n\n")
# summary of model after adding new params
model.set_action_normalizer(actions_normalizer)
model.set_state_normalizer(observations_normalizer)
print("\n\n -----summary of model AFTER adding normalization modules")
print("num trainable parameters", count_trainable_parameters(model))
print("num non-trainable parameters ",
count_non_trainable_parameters(model))
print("\n\n")
# unsqueeze to mimic dataloader with batch size of 1
data = dataset[0] # test the getitem
observations = data['observations'].unsqueeze(0)
actions = data['actions'].unsqueeze(0)
obs_slice = observations[:, :n_history, :]
action_slice = actions[:, :n_history, :]
print("action_slice.shape", action_slice.shape)
print("obs_slice.shape", obs_slice.shape)
# run the model forwards one timestep
output = model.forward(obs_slice, action_slice)
print("output.shape", output.shape)
# save the model with torch.save and torch.load
save_dir = os.path.join(get_project_root(), 'sandbox')
model_save_file = os.path.join(save_dir, "model.pth")
torch.save(model, model_save_file)
# load the model
model_load = torch.load(model_save_file)
print("\n\n -----summary of model LOADED from disk")
print("num trainable parameters", count_trainable_parameters(model_load))
print("num non-trainable parameters ",
count_non_trainable_parameters(model_load))
print("\n\n")
# now try doing the same but with the state dict
# my hunch is that this won't work . . .
params_save_file = os.path.join(save_dir, "model_params.pth")
torch.save(model.state_dict(), params_save_file)
# load the model
model_load = DynaNetMLP(config)
state_dict = torch.load(params_save_file)
for param_tensor in state_dict:
print(param_tensor, "\t", state_dict[param_tensor].size())
# try creating some dummy DataNormalizer objects
# model_load.set_state_normalizer(DataNormalizer(0.0,1.0))
# model_load.set_action_normalizer(DataNormalizer(0.0,1.0))
model_load.load_state_dict(state_dict)
print("\n\n -----summary of model LOADED from disk with state_dict method")
print("num trainable parameters", count_trainable_parameters(model_load))
print("num non-trainable parameters ",
count_non_trainable_parameters(model_load))
print("\n\n")
print("model_load._action_normalizer._mean",
model_load.action_normalizer._mean)
print("model._action_normalizer._mean", model.action_normalizer._mean)
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']))
import os
from key_dynam.utils.utils import get_project_root, get_data_root
SIM_ASSETS_ROOT = os.path.join(get_project_root(), 'sim_assets')
block_push = os.path.join(SIM_ASSETS_ROOT, 'block_push.urdf')
extra_heavy_duty_table = os.path.join(SIM_ASSETS_ROOT, "extra_heavy_duty_table_surface_only_collision.sdf")
xy_slide = os.path.join(SIM_ASSETS_ROOT, "xy_slide.urdf")
ycb_model_paths = dict({
'cracker_box': os.path.join(SIM_ASSETS_ROOT, "cracker_box/003_cracker_box.sdf"),
'sugar_box': os.path.join(SIM_ASSETS_ROOT, "sugar_box/004_sugar_box.sdf"),
'tomato_soup_can': os.path.join(SIM_ASSETS_ROOT, "tomato_soup_can/005_tomato_soup_can.sdf"),
'mustard_bottle': os.path.join(SIM_ASSETS_ROOT, "mustard_bottle/006_mustard_bottle.sdf"),
'gelatin_box': os.path.join(SIM_ASSETS_ROOT, "gelatin_box/009_gelatin_box.sdf"),
'potted_meat_can': os.path.join(SIM_ASSETS_ROOT, "potted_meat_can/010_potted_meat_can.sdf")
})
ycb_model_baselink_names = dict({
'cracker_box': 'base_link_cracker',
'sugar_box': 'base_link_sugar',
'tomato_soup_can': 'base_link_soup',
'mustard_bottle': 'base_link_mustard',
'gelatin_box': 'base_link_gelatin',
'potted_meat_can': 'base_link_meat'
})
LARGE_SIM_ASSETS_ROOT = os.path.join(get_data_root(), 'stable/sim_assets')
from key_dynam.utils import utils, torch_utils
from key_dynam.controller.zmq_utils import ZMQServer
from key_dynam.dataset.online_episode_reader import OnlineEpisodeReader
from key_dynam.dataset.mpc_dataset import DynamicsModelInputBuilder
from key_dynam.controller.plan_container import PlanContainer
from key_dynam.dynamics.models_dy import rollout_model, rollout_action_sequences, get_object_and_robot_state_indices
from key_dynam.planner import utils as planner_utils
from key_dynam.utils import meshcat_utils
from key_dynam.experiments.exp_22_push_box_hardware.visualize_dynamics_model import \
visualize_model_prediction_single_timestep
# dense correspondence
from dense_correspondence_manipulation.utils import meshcat_utils as pdc_meshcat_utils
from dense_correspondence_manipulation.utils.constants import DEPTH_IM_SCALE
PLAN_MSG_FILE = os.path.join(utils.get_project_root(), 'sandbox', 'plan_msg.p')
COMPUTE_CONTROL_ACTION_MSG_FILE = os.path.join(utils.get_project_root(),
'sandbox',
'compute_control_action_msg.p')
SAVE_MESSAGES = True
class ControllerState(Enum):
STOPPED = 0
PLAN_READY = 1
RUNNING = 2
class Controller:
"""
def get_env_config():
return load_yaml(
os.path.join(get_project_root(),
'experiments/exp_29_mugs/config.yaml'))
CUDA_VISIBLE_DEVICES = [0] set_cuda_visible_devices(CUDA_VISIBLE_DEVICES) config = dict() config['train'] = dict() config['train']['train_valid_ratio'] = 0.9 DATASET_NAME = "top_down_rotated" # DATASET_NAME = "2019-12-05-15-58-48-462834_top_down_rotated_250" # # MODEL_NAME = "2019-12-04-01-32-12-010393_top_down_rotated_sigma_5" # MODEL_NAME = "2019-12-26-20-49-35-944397_dataset_2019-12-05-15-58-48-462834_top_down_rotated_250" MODEL_NAME = "2019-12-27-22-29-23-359200_dataset_2019-12-05-15-58-48-462834_top_down_rotated_250" dataset_root = os.path.join(get_project_root(), "data/dev/experiments/05/data", DATASET_NAME) multi_episode_dict = DrakeSimEpisodeReader.load_dataset(dataset_root) network_folder = os.path.join(get_project_root(), "data/dev/experiments/05/trained_models", MODEL_NAME) epoch = 55 iter = 0 model_file = os.path.join(network_folder, "net_dy_epoch_%d_iter_%d_model.pth" %(epoch, iter)) model = torch.load(model_file) model.cuda() model = model.eval() # make this unique output_dir = os.path.join(get_project_root(), "data/dev/experiments/05/precomputed_descriptor_keypoints", DATASET_NAME, MODEL_NAME, get_current_YYYY_MM_DD_hh_mm_ss_ms())
import os
import pydrake
import time
from dense_correspondence_manipulation.utils.utils import set_cuda_visible_devices
set_cuda_visible_devices([1])
from key_dynam.utils.utils import get_project_root, save_yaml, load_yaml, get_current_YYYY_MM_DD_hh_mm_ss_ms, get_data_root
from key_dynam.transporter.train_transporter import train_transporter
from key_dynam.dense_correspondence.dc_drake_sim_episode_reader import DCDrakeSimEpisodeReader
# project root
PROJECT_ROOT = get_project_root()
# load config
config = load_yaml(
os.path.join(get_project_root(),
"experiments/exp_18_box_on_side/config.yaml"))
# specify the dataset
dataset_name = "dps_box_on_side_600"
config['perception']['dataset_name'] = dataset_name
dataset_name = config['perception']['dataset_name']
config['dataset']['set_epoch_size_to_num_images'] = False
config['dataset']['epoch_size'] = {'train': 6000, 'valid': 300}
# camera_name = "camera_angled"
camera_name = "camera_1_top_down"
def main():
# load dynamics model
model_dict = load_model_state_dict()
model = model_dict['model_dy']
model_dd = model_dict['model_dd']
config = model.config
env_config = load_yaml(os.path.join(get_project_root(), 'experiments/exp_20_mugs/config.yaml'))
env_config['env']['observation']['depth_int16'] = True
n_history = config['train']['n_history']
initial_cond = generate_initial_condition(env_config, push_length=PUSH_LENGTH)
env_config = initial_cond['config']
# enable the right observations
camera_name = model_dict['metadata']['camera_name']
spatial_descriptor_data = model_dict['spatial_descriptor_data']
ref_descriptors = spatial_descriptor_data['spatial_descriptors']
K = ref_descriptors.shape[0]
ref_descriptors = torch.Tensor(ref_descriptors).cuda() # put them on the GPU
print("ref_descriptors\n", ref_descriptors)
print("ref_descriptors.shape", ref_descriptors.shape)
# create the environment
# create the environment
env = DrakeMugsEnv(env_config)
env.reset()
T_world_camera = env.camera_pose(camera_name)
camera_K_matrix = env.camera_K_matrix(camera_name)
# create another environment for doing rollouts
env2 = DrakeMugsEnv(env_config, visualize=False)
env2.reset()
action_function = ActionFunctionFactory.function_from_config(config)
observation_function = ObservationFunctionFactory.drake_pusher_position_3D(config)
visual_observation_function = \
VisualObservationFunctionFactory.descriptor_keypoints_3D(config=config,
camera_name=camera_name,
model_dd=model_dd,
ref_descriptors=ref_descriptors,
K_matrix=camera_K_matrix,
T_world_camera=T_world_camera,
)
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()
reset_environment(env, initial_cond['q_pusher'], initial_cond['q_slider'])
obs_init = env.get_observation()
#### 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)
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()
#
idx = episode.get_latest_idx()
obs_raw = episode.get_observation(idx)
z_object_goal = goal_func(obs_raw)
z_keypoints_init_W = keypoints_3D_from_dynamics_model_output(z_object_goal, K)
z_keypoints_init_W = torch_utils.cast_to_numpy(z_keypoints_init_W)
z_keypoints_obj = keypoints_world_frame_to_object_frame(z_keypoints_init_W,
T_W_obj=slider_pose_from_observation(obs_init))
color = [1, 0, 0]
meshcat_utils.visualize_points(vis=vis,
name="keypoints_W",
pts=z_keypoints_init_W,
color=color,
size=0.02,
)
# input("press Enter to continue")
# rollout single action sequence using the simulator
action_sequence_np = torch_utils.cast_to_numpy(initial_cond['action_sequence'])
N = action_sequence_np.shape[0]
obs_rollout_gt = env_utils.rollout_action_sequence(env, action_sequence_np)[
'observations']
# using the vision model to get "goal" keypoints
z_object_goal = goal_func(obs_rollout_gt[-1])
z_object_goal_np = torch_utils.cast_to_numpy(z_object_goal)
z_keypoints_goal = keypoints_3D_from_dynamics_model_output(z_object_goal, K)
z_keypoints_goal = torch_utils.cast_to_numpy(z_keypoints_goal)
# visualize goal keypoints
color = [0, 1, 0]
meshcat_utils.visualize_points(vis=vis,
name="goal_keypoints",
pts=z_keypoints_goal,
color=color,
size=0.02,
)
# 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)
action_seq_gt_torch = torch_utils.cast_to_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']
print("state_pred_rollout.shape", state_pred_rollout.shape)
for i in range(N):
# vis GT for now
name = "GT_3D/%d" % (i)
T_W_obj = slider_pose_from_observation(obs_rollout_gt[i])
# print("T_W_obj", T_W_obj)
# green
color = np.array([0, 1, 0]) * get_color_intensity(i, N)
meshcat_utils.visualize_points(vis=vis,
name=name,
pts=z_keypoints_obj,
color=color,
size=0.01,
T=T_W_obj)
# red
color = np.array([0, 0, 1]) * get_color_intensity(i, N)
state_pred = state_pred_rollout[:, i, :]
pts_pred = keypoints_3D_from_dynamics_model_output(state_pred, K).squeeze()
pts_pred = pts_pred.detach().cpu().numpy()
name = "pred_3D/%d" % (i)
meshcat_utils.visualize_points(vis=vis,
name=name,
pts=pts_pred,
color=color,
size=0.01,
)
# input("finished visualizing GT rollout\npress Enter to continue")
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 = 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
if n_look_ahead == 0:
break
# 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 mpc_out is not None:
action_seq_rollout_init = mpc_out['action_seq'][1:]
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)
mpc_out = planner.trajectory_optimization(state_cur=z_cur,
action_his=action_his,
obs_goal=z_object_goal_np,
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,
)
print("MPC step took %.4f seconds" %(time.time() - mpc_start_time))
action_seq_mpc = mpc_out['action_seq'].cpu().numpy()
# Rollout with ground truth simulator dynamics
action_seq_mpc = torch_utils.cast_to_numpy(mpc_out['action_seq'])
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']
state_pred_mpc = torch_utils.cast_to_numpy(mpc_out['state_pred'])
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):
# red
color = np.array([1, 0, 0]) * get_color_intensity(i, L)
state_pred = state_pred_mpc[i, :]
state_pred = np.expand_dims(state_pred, 0) # may need to expand dims here
pts_pred = keypoints_3D_from_dynamics_model_output(state_pred, K).squeeze()
name = "mpc_3D/%d" % (i)
meshcat_utils.visualize_points(vis=vis,
name=name,
pts=pts_pred,
color=color,
size=0.01,
)
# ground truth rollout of the MPC action_seq
name = "mpc_GT_3D/%d" % (i)
T_W_obj = slider_pose_from_observation(obs_mpc_gt[i])
# green
color = np.array([1, 1, 0]) * get_color_intensity(i, L)
meshcat_utils.visualize_points(vis=vis,
name=name,
pts=z_keypoints_obj,
color=color,
size=0.01,
T=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)
# visualize current keypoint positions
obs_cur = env.get_observation()
T_W_obj = slider_pose_from_observation(obs_cur)
# yellow
color = np.array([1, 1, 0])
meshcat_utils.visualize_points(vis=vis,
name="keypoint_cur",
pts=z_keypoints_obj,
color=color,
size=0.02,
T=T_W_obj)
action_seq_mpc = action_seq_mpc[1:]
state_pred_mpc = state_pred_mpc[1:]
obs_final = env.get_observation()
pose_error = compute_pose_error(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 load_model_and_data():
dataset_name = "push_box_hardware"
# model_name = "DD_2D/2020-06-24-22-22-58-234812_DD_3D_n_his_2" # this model is actually 3D
# model_name = "DD_3D/2020-06-25-00-49-29-679042_DD_3D_n_his_2_T_aug"
# model_name = "DD_3D/2020-06-25-00-39-29-020621_DD_3D_n_his_2"
model_name = "DD_3D/2020-07-02-17-59-21-362337_DD_3D_n_his_2_T_aug"
train_dir = "/home/manuelli/data/key_dynam/dev/experiments/22/dataset_push_box_hardware/trained_models/dynamics"
train_dir = os.path.join(train_dir, model_name)
ckpt_file = os.path.join(train_dir, "net_best_dy_state_dict.pth")
config = load_yaml(os.path.join(train_dir, 'config.yaml'))
state_dict = torch.load(ckpt_file)
# build dynamics model
model_dy = build_dynamics_model(config)
# print("state_dict.keys()", state_dict.keys())
model_dy.load_state_dict(state_dict)
model_dy = model_dy.eval()
model_dy = model_dy.cuda()
spatial_descriptor_data = load_pickle(
os.path.join(train_dir, 'spatial_descriptors.p'))
metadata = load_pickle(os.path.join(train_dir, 'metadata.p'))
# build dense-descriptor model
model_dd_file = metadata['model_file']
model_dd = torch.load(model_dd_file)
model_dd = model_dd.eval()
model_dd = model_dd.cuda()
# load the dataset
dataset_paths = get_dataset_paths(dataset_name)
dataset_root = dataset_paths['dataset_root']
episodes_config = dataset_paths['episodes_config']
precomputed_vision_data_root = DD_utils.get_precomputed_data_root(
dataset_name)['precomputed_data_root']
# descriptor_keypoints_root = os.path.join(precomputed_vision_data_root, 'descriptor_keypoints')
descriptor_keypoints_root = os.path.join(precomputed_vision_data_root,
'descriptor_keypoints')
multi_episode_dict = DynamicSpartanEpisodeReader.load_dataset(
config=config,
episodes_config=episodes_config,
episodes_root=dataset_paths['dataset_root'],
load_image_episode=True,
precomputed_data_root=descriptor_keypoints_root,
max_num_episodes=None)
visual_observation_function = PrecomputedVisualObservationFunctionFactory.function_from_config(
config,
keypoint_idx=spatial_descriptor_data['spatial_descriptors_idx'])
action_function = ActionFunctionFactory.function_from_config(config)
observation_function = ObservationFunctionFactory.function_from_config(
config)
dataset = MultiEpisodeDataset(
config,
action_function=action_function,
observation_function=observation_function,
episodes=multi_episode_dict,
visual_observation_function=visual_observation_function,
phase="valid", # this means no data augmentation
)
#### PLANNER #######
planner = None
# make a planner config
planner_config = copy.copy(model_dy.config)
config_tmp = load_yaml(
os.path.join(get_project_root(),
'experiments/exp_22_push_box_hardware/config_DD_3D.yaml'))
planner_config['mpc'] = config_tmp['mpc']
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))
return {
"model_dy": model_dy,
'model_dd': model_dd,
'dataset': dataset,
'config': config,
"multi_episode_dict": multi_episode_dict,
'spatial_descriptor_data': spatial_descriptor_data,
'planner': planner,
'observation_function': observation_function,
'action_function': action_function,
}
