def collect_episodes(config, output_dir=None, visualize=True):
if output_dir is None:
output_dir = os.path.join(os.getcwd(), 'data')
if not os.path.exists(output_dir):
os.makedirs(output_dir)
# save the config
config_save_file = os.path.join(output_dir, 'config.yaml')
save_yaml(config, config_save_file)
save_file = os.path.join(output_dir,
"%s.p" % (get_current_YYYY_MM_DD_hh_mm_ss_ms()))
# initialize config for DataCollector
dc = PusherSliderDataCollector(config)
num_episodes = config['dataset']['num_episodes']
multi_episode_container = MultiEpisodeContainer()
for i in range(num_episodes):
print("collecting episode %d of %d" % (i + 1, num_episodes))
name = "%s_idx_%d" % (get_current_YYYY_MM_DD_hh_mm_ss_ms(), i)
episode = dc.collect_single_episode(visualize, episode_name=name)
multi_episode_container.add_episode(episode)
print("saving data to %s" % save_file)
multi_episode_container.save_to_file(save_file)
def train_dense_descriptor_vision(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
dataset_paths = get_dataset_paths(dataset_name)
episodes_config = dataset_paths['episodes_config']
episodes_root = dataset_paths['dataset_root']
# data augmentation
if True:
model_name = "data_aug_%s" % (get_current_YYYY_MM_DD_hh_mm_ss_ms())
output_dir = os.path.join(
get_experiment_save_root(dataset_name),
'trained_models/perception/dense_descriptors', model_name)
dataset_paths = get_dataset_paths(dataset_name)
dataset_root = dataset_paths['dataset_root']
multi_episode_dict = DCDynamicSpartanEpisodeReader.load_dataset(
config=episodes_config, episodes_root=episodes_root)
config = load_config()
config['dataset']['data_augmentation'] = True
config['dataset']['camera_names'] = dataset_paths[
'dense_descriptor_camera_list']
train_dense_descriptors(config,
train_dir=output_dir,
multi_episode_dict=multi_episode_dict,
verbose=False)
def save_data(
self,
save_dir=None,
):
"""
Saves data from the
- PlanContainer
- OnlineEpisode
"""
if save_dir is None:
save_dir = os.path.join(
utils.get_data_root(),
'hardware_experiments/closed_loop_rollouts/sandbox',
utils.get_current_YYYY_MM_DD_hh_mm_ss_ms())
if not os.path.exists(save_dir):
os.makedirs(save_dir)
print("saving MPC rollout data at: %s" % (save_dir))
save_data = {
'episode': self._state_dict['episode'].get_save_data(),
'plan': self._state_dict['plan'].get_save_data(),
}
save_file = os.path.join(save_dir, "data.p")
utils.save_pickle(save_data, save_file)
print("done saving data")
def __init__(self, name=None):
if name is None:
name = get_current_YYYY_MM_DD_hh_mm_ss_ms()
self._data = dict()
self._data['name'] = name
self._data['config'] = None
self._data['trajectory'] = []
def collect_episodes(config,
output_dir=None,
visualize=True,
use_threads=False):
if output_dir is None:
output_dir = os.path.join(os.getcwd(), 'data')
if not os.path.exists(output_dir):
os.makedirs(output_dir)
# save the config
config_save_file = os.path.join(output_dir, 'config.yaml')
save_yaml(config, config_save_file)
# initialize config for DataCollector
dc = DrakePusherSliderEpisodeCollector(config)
num_episodes = config['dataset']['num_episodes']
# record some metadata
metadata = dict()
metadata['episodes'] = dict()
for i in range(num_episodes):
print("\n")
start_time = time.time()
print("collecting episode %d of %d" % (i + 1, num_episodes))
name = "%s_idx_%d" % (get_current_YYYY_MM_DD_hh_mm_ss_ms(), i)
episode = dc.collect_single_episode(visualize, episode_name=name)
print("saving to disk")
metadata['episodes'][name] = dict()
image_data_file = episode.save_images_to_hdf5(output_dir)
non_image_data_file = episode.save_non_image_data_to_pickle(output_dir)
print("non_image_data.keys()", episode.non_image_data.keys())
metadata['episodes'][name]['non_image_data_file'] = non_image_data_file
metadata['episodes'][name]['image_data_file'] = image_data_file
print("done saving to disk")
elapsed = time.time() - start_time
print("single episode took: %.2f seconds" % (elapsed))
save_yaml(metadata, os.path.join(output_dir, 'metadata.yaml'))
def collect_single_episode(
env, # gym env
action_seq, # [N, action_dim]
episode_name=None):
"""
Rollout the action sequence using the simulator
Record the observations
"""
if episode_name is None:
episode_name = get_current_YYYY_MM_DD_hh_mm_ss_ms()
episode_container = EpisodeContainer()
episode_container.set_config(env.config)
episode_container.set_name(episode_name)
episode_container.set_metadata(env.get_metadata())
obs_list = []
N = action_seq.shape[0]
obs_prev = env.get_observation()
for i in range(N):
action = action_seq[i]
obs, reward, done, info = env.step(action)
episode_container.add_obs_action(obs_prev, action)
obs_prev = obs
# terminate if outside boundary
if not env.slider_within_boundary():
print('slider outside boundary, terminating episode')
break
if not env.pusher_within_boundary():
print('pusher ouside boundary, terminating episode')
break
return {'episode_container': episode_container}
def collect_episodes(config,
output_dir=None,
visualize=True,
debug=False):
if output_dir is None:
output_dir = os.path.join(os.getcwd(), 'data')
if not os.path.exists(output_dir):
os.makedirs(output_dir)
# save the config
config_save_file = os.path.join(output_dir, 'config.yaml')
save_yaml(config, config_save_file)
# initialize config for DataCollector
dc = DrakePusherSliderEpisodeCollector(config, visualize=visualize)
num_episodes = config['dataset']['num_episodes']
# record some metadata
metadata = dict()
metadata['episodes'] = dict()
while (len(metadata['episodes']) < num_episodes):
i = len(metadata['episodes'])
if debug:
input("Press Enter to continue...")
print("\n")
start_time = time.time()
print("collecting episode %d of %d" % (i + 1, num_episodes))
name = "%s_idx_%d" % (get_current_YYYY_MM_DD_hh_mm_ss_ms(), i)
ic = generate_initial_condition(config=config,
T_aug_enabled=True)
if debug:
print("initial condition\n", ic)
episode = dc.collect_single_episode(visualize=visualize,
episode_name=name,
q_pusher=ic['q_pusher'],
q_slider=ic['q_slider'],
# v_pusher=ic['v_pusher'],
action_seq=ic['action_seq'],
)
# potentially discard it if the object didn't move during the data collection
if len(episode._data['trajectory']) < 10:
print("trajectory was too short, skipping")
continue
obs_start = episode._data['trajectory'][5]['observation']
obs_end = episode._data['trajectory'][-1]['observation']
q_slider_start = obs_start['slider']['position']['translation']
q_slider_end = obs_end['slider']['position']['translation']
dq_slider = obs_start['slider']['position']['translation'] - obs_end['slider']['position']['translation']
if debug:
print("len(episode._data['trajectory'])", len(episode._data['trajectory']))
print("q_slider_start", q_slider_start)
print("q_slider_end", q_slider_end)
print("dq_slider", dq_slider)
print("np.linalg.norm(dq_slider)", np.linalg.norm(dq_slider))
# if slider didn't move by at least 1 mm then discard this episode
if (np.linalg.norm(dq_slider) < 0.001): # one mm
print("discarding episode since slider didn't move")
continue
print("saving to disk")
metadata['episodes'][name] = dict()
image_data_file = episode.save_images_to_hdf5(output_dir)
non_image_data_file = episode.save_non_image_data_to_pickle(output_dir)
print("output_dir:", output_dir)
print("non_image_data.keys()", episode.non_image_data.keys())
metadata['episodes'][name]['non_image_data_file'] = non_image_data_file
metadata['episodes'][name]['image_data_file'] = image_data_file
print("done saving to disk")
elapsed = time.time() - start_time
print("single episode took: %.2f seconds" % (elapsed))
save_yaml(metadata, os.path.join(output_dir, 'metadata.yaml'))
def main():
dataset_name, multi_episode_dict = load_episodes()
## Load Model
model_name, model_file = get_DD_model_file()
model = torch.load(model_file)
model = model.cuda()
model = model.eval()
# make this unique
output_dir = os.path.join(get_data_root(),
"dev/experiments/09/precomputed_vision_data",
"dataset_%s" % (dataset_name),
"model_name_%s" % (model_name),
get_current_YYYY_MM_DD_hh_mm_ss_ms())
camera_name = "camera_1_top_down"
episode_name = "2020-05-13-21-55-01-487901_idx_33"
episode_idx = 22
# compute descriptor confidence scores
if True:
print("\n\n---------Computing Descriptor Confidence Scores-----------")
metadata_file = os.path.join(output_dir, 'metadata.p')
if os.path.isfile(metadata_file):
answer = input(
"metadata.p file already exists, do you want to overwrite it? y/n"
)
if answer == "y":
os.rmdir(output_dir)
print("removing existing file and continuing")
else:
print("aborting")
quit()
compute_descriptor_confidences(
multi_episode_dict,
model,
output_dir,
batch_size=10,
num_workers=20,
model_file=model_file,
camera_name=camera_name,
num_ref_descriptors=50,
num_batches=10,
episode_name_arg=episode_name,
episode_idx=episode_idx,
)
if True:
metadata_file = os.path.join(output_dir, 'metadata.p')
metadata = load_pickle(metadata_file)
# metadata_file = "/media/hdd/data/key_dynam/dev/experiments/09/precomputed_vision_data/dataset_2020-03-25-19-57-26-556093_constant_velocity_500/model_name_2020-04-07-14-31-35-804270_T_aug_dataset/2020-04-09-20-51-50-624799/metadata.p"
# metadata = load_pickle(metadata_file)
print("\n\n---------Precomputing Descriptor Keypoints-----------")
descriptor_keypoints_output_dir = os.path.join(output_dir,
"descriptor_keypoints")
precompute_descriptor_keypoints(
multi_episode_dict,
model,
descriptor_keypoints_output_dir,
ref_descriptors_metadata=metadata,
batch_size=10,
num_workers=20,
)
if True:
confidence_score_data_file = os.path.join(output_dir, 'data.p')
confidence_score_data = load_pickle(confidence_score_data_file)
print(
"\n\n---------Selecting Spatially Separated Keypoints-----------")
score_and_select_spatially_separated_keypoints(
metadata,
confidence_score_data=confidence_score_data,
K=5,
position_diff_threshold=25,
output_dir=output_dir,
)
print("Data saved at: ", output_dir)
print("Finished Normally")
def collect_single_episode(
self,
visualize=True,
episode_name=None,
q_pusher=None,
q_slider=None,
v_pusher=None,
action_sampler=None, # functional
action_seq=None, # sequence of actions [N, action_dim]
):
"""
This collects a single episode by performing the following steps
- sample initial conditions for environment
- sample action for environment
- collect the interaction
:return: EpisodeContainer
:rtype:
"""
if episode_name is None:
episode_name = get_current_YYYY_MM_DD_hh_mm_ss_ms()
env = self._env
env.reset()
if visualize:
print("setting target realtime rate 1.0")
env.simulator.set_target_realtime_rate(1.0)
# it takes a while for this to take effect, so it might initially
# look like the sim is running really fast, but this is
# just a cosmetic issue
# sample if they weren't passed in
if q_pusher is None or q_slider is None:
q_pusher, q_slider = self.get_pusher_slider_initial_positions()
use_constant_action = (v_pusher is None)
if action_sampler is None:
action_sampler = self.sample_pusher_velocity
def get_action(counter):
if v_pusher is not None:
return v_pusher
elif action_seq is not None:
return action_seq[counter]
elif action_sampler is not None:
return action_sampler()
else:
return self.sample_pusher_velocity()
context = env.get_mutable_context()
env.set_pusher_position(context, q_pusher)
env.set_slider_position(context, q_slider)
action_zero = np.zeros(2)
# if visualize:
# env.step(action_zero)
# print("sleeping for several seconds to let meshcat visualize")
# time.sleep(5.0)
# log data into EpisodeContainer object
episode_container = EpisodeContainer()
episode_container.set_config(env.config)
episode_container.set_name(episode_name)
episode_container.set_metadata(env.get_metadata())
# step for 2 seconds to let object fall
env.step(action_zero, dt=2.0)
obs_prev = env.get_observation()
start_time = time.time()
num_timesteps = self._config['dataset'][
'num_timesteps'] + self._config['train']['n_history']
counter = 0
for i in range(num_timesteps):
start_time_1 = time.time()
# stand still for n_history timesteps
if i < self._config['train']['n_history']:
action = action_zero # just sit still
# this makes it easier to train our dynamics model
# since we need to get a history of observations
else:
action = get_action(counter)
counter += 1
# exploration_type = self._config['dataset']['data_generation']['exploration_type']
# if exploration_type == 'random':
# action = self.sample_pusher_velocity()
# elif exploration_type == 'constant':
# action = action_actual
# single sim time
obs, reward, done, info = env.step(action)
episode_container.add_obs_action(obs_prev, action)
obs_prev = obs
if visualize:
# note that sim wil
elapsed = time.time() - start_time_1
hz = 1.0 / elapsed
# print("sim rate (hz)", hz)
# print("pusher velocity", obs['pusher']['velocity'])
# terminate if outside boundary
if not env.slider_within_boundary():
print('slider outside boundary, terminating episode')
break
if not env.pusher_within_boundary():
print('pusher ouside boundary, terminating episode')
break
elapsed = time.time() - start_time
if DEBUG:
print("elapsed wall clock", elapsed)
print("elapsed sim time", obs_prev['sim_time'])
return episode_container
def DD_3D_dynamics(dataset_name):
from key_dynam.training.train_dynamics_pusher_slider_precomputed_keypoints import train_dynamics
from key_dynam.experiments.drake_pusher_slider import DD_utils
def load_config():
return load_yaml(
os.path.join(
get_project_root(),
'experiments/exp_22_push_box_hardware/config_DD_3D.yaml'))
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')
config = load_config()
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)
experiment_save_root = get_experiment_save_root(dataset_name)
# experiment_save_root = os.path.join(get_data_root(), 'sandbox')
# standard
if True:
TRAIN = True
spatial_descriptor_data = load_pickle(
os.path.join(precomputed_vision_data_root,
'spatial_descriptors.p'))
metadata = load_pickle(
os.path.join(precomputed_vision_data_root, 'metadata.p'))
train_dir = None
config = load_config()
config['train']['n_history'] = 2
config['train']['random_seed'] = 1
config['dataset']['visual_observation_function'][
'camera_name'] = dataset_paths['main_camera_name']
config['dataset'][
'precomputed_data_root'] = precomputed_vision_data_root
suffix = "_DD_3D_n_his_2"
model_name = get_current_YYYY_MM_DD_hh_mm_ss_ms() + suffix
if TRAIN:
train_dir = os.path.join(experiment_save_root,
'trained_models/dynamics/DD_3D/',
model_name)
os.makedirs(train_dir)
visual_observation_function = PrecomputedVisualObservationFunctionFactory.function_from_config(
config,
keypoint_idx=spatial_descriptor_data['spatial_descriptors_idx']
)
train_dynamics(
config=config,
train_dir=train_dir,
multi_episode_dict=multi_episode_dict,
visual_observation_function=visual_observation_function,
metadata=metadata,
spatial_descriptors_data=spatial_descriptor_data)
# data_aug
if False:
TRAIN = True
spatial_descriptor_data = load_pickle(
os.path.join(precomputed_vision_data_root,
'spatial_descriptors.p'))
metadata = load_pickle(
os.path.join(precomputed_vision_data_root, 'metadata.p'))
train_dir = None
config = load_config()
config['train']['n_history'] = 2
config['train']['random_seed'] = 1
config['dataset']['visual_observation_function'][
'camera_name'] = dataset_paths['main_camera_name']
config['dataset']['data_augmentation']['enabled'] = True
config['dataset'][
'precomputed_data_root'] = precomputed_vision_data_root
suffix = "_DD_3D_n_his_2_T_aug"
model_name = get_current_YYYY_MM_DD_hh_mm_ss_ms() + suffix
if TRAIN:
train_dir = os.path.join(experiment_save_root,
'trained_models/dynamics/DD_3D/',
model_name)
os.makedirs(train_dir)
visual_observation_function = PrecomputedVisualObservationFunctionFactory.function_from_config(
config,
keypoint_idx=spatial_descriptor_data['spatial_descriptors_idx']
)
train_dynamics(
config=config,
train_dir=train_dir,
multi_episode_dict=multi_episode_dict,
visual_observation_function=visual_observation_function,
metadata=metadata,
spatial_descriptors_data=spatial_descriptor_data)
# load config
# specify the dataset
dataset_name = "dps_box_on_side_600"
# prepare folders
data_dir = os.path.join(get_data_root(), 'dev/experiments/18/data',
dataset_name)
start_time = time.time()
multi_episode_dict = DCDrakeSimEpisodeReader.load_dataset(
dataset_root=data_dir)
# placeholder for now
config = dev_utils.load_integral_heatmap_3d_config()
config['dataset']['name'] = dataset_name
config['dataset']['camera_names'] = ['camera_angled', 'camera_angled_rotated']
config['dataset']['train_valid_ratio'] = 0.83
model_name = "3D_loss_%s_%s" % ("camera_angled",
get_current_YYYY_MM_DD_hh_mm_ss_ms())
train_dir = os.path.join(
get_data_root(),
'dev/experiments/18/trained_models/perception/dense_descriptor',
model_name)
train_dense_descriptors(config,
train_dir=train_dir,
multi_episode_dict=multi_episode_dict,
verbose=False)
# camera_name = "camera_angled"
camera_name = "camera_1_top_down"
config['perception']['camera_name'] = camera_name
config['train'] = config['train_transporter']
config['train']['train_valid_ratio'] = 0.83
print(config)
# prepare folders
data_dir = os.path.join(get_data_root(), 'dev/experiments/18/data',
dataset_name)
train_dir = os.path.join(
get_data_root(), 'dev/experiments/18/trained_models/perception',
"transporter_%s_%s" % (camera_name, get_current_YYYY_MM_DD_hh_mm_ss_ms()))
ckp_dir = os.path.join(
train_dir, 'train_nKp%d_invStd%.1f' %
(config['perception']['n_kp'], config['perception']['inv_std']))
start_time = time.time()
multi_episode_dict = DCDrakeSimEpisodeReader.load_dataset(
dataset_root=data_dir)
print("loading dataset took %d seconds", time.time() - start_time)
train_transporter(config=config,
train_dir=train_dir,
ckp_dir=ckp_dir,
multi_episode_dict=multi_episode_dict)
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())
camera_name = "camera_1_top_down"
precompute_descriptor_keypoints(multi_episode_dict,
model,
output_dir,
batch_size=10,
num_workers=20,
model_file=model_file,
camera_name=camera_name,
num_ref_descriptors=32,
)
def collect_episodes(config,
metadata,
data_collector,
n_episode=1,
data_dir=None,
visualize=False,
episode_name=None,
exploration_type='random',
model_dy=None):
"""
This collects a single episode by performing the following steps
- sample initial conditions for environment
- sample action from controller for environment
- collect the interaction
"""
if data_dir is None:
data_dir = os.path.join(os.getcwd(), 'data')
if not os.path.exists(data_dir):
os.makedirs(data_dir)
for idx_episode in range(n_episode):
start_time = time.time()
print("collecting episode %d of %d" % (idx_episode + 1, n_episode))
episode_name = "%s_idx_%d" % (get_current_YYYY_MM_DD_hh_mm_ss_ms(),
idx_episode)
env = data_collector._env
env.reset()
if visualize:
print("setting target realtime rate 1.0")
env.simulator.set_target_realtime_rate(1.0)
q_pusher, q_slider = data_collector.get_pusher_slider_initial_positions(
)
context = env.get_mutable_context()
env.set_pusher_position(context, q_pusher)
env.set_slider_position(context, q_slider)
action_zero = np.zeros(2)
env.step(action_zero, dt=1.0) # to get the box to drop down
# log data into EpisodeContainer object
episode_container = EpisodeContainer()
episode_container.set_config(env.config)
episode_container.set_name(episode_name)
episode_container.set_metadata(env.get_metadata())
obs_prev = env.get_observation()
num_timesteps = config['mpc']['num_timesteps'] + config['train'][
'n_history']
for i in range(num_timesteps):
# stand still for n_history timesteps
if i < config['train']['n_history']:
action = action_zero
# this makes it easier to train our dynamics model
# since we need to get a history of observations
else:
if exploration_type == 'random':
action = data_collector.sample_pusher_velocity()
elif exploration_type == 'mppi':
action = data_collector.sample_pusher_velocity()
# single sim time
obs, reward, done, info = env.step(action)
episode_container.add_obs_action(obs_prev, action)
obs_prev = obs
# terminate if outside boundary
if not env.slider_within_boundary():
print('slider outside boundary, terminating episode')
break
if not env.pusher_within_boundary():
print('pusher ouside boundary, terminating episode')
break
# print("saving to disk")
metadata['episodes'][episode_name] = dict()
image_data_file = episode_container.save_images_to_hdf5(data_dir)
non_image_data_file = episode_container.save_non_image_data_to_pickle(
data_dir)
metadata['episodes'][episode_name][
'non_image_data_file'] = non_image_data_file
metadata['episodes'][episode_name]['image_data_file'] = image_data_file
# print("done saving to disk")
elapsed = time.time() - start_time
print("single episode took: %.2f seconds" % (elapsed))
def collect_single_episode(self, visualize=True, episode_name=None):
"""
This collects a single episode by performing the following steps
- sample initial conditions for environment
- sample action for environment
- collect the interaction
:return: EpisodeContainer
:rtype:
"""
if episode_name is None:
episode_name = get_current_YYYY_MM_DD_hh_mm_ss_ms()
env = self._env
env.reset()
# set pusher position
# this is in Body frame
pusher_position_B = self.pusher_initial_position_space_from_config(
env.config).sample()
# this is in world frame
pusher_position_W = env.slider_body.local_to_world(pusher_position_B)
# set pusher position
env.pusher_body.position = pusher_position_W
# step so we can get an observation that has correct image
env.step([0, 0])
# get action
# don't worry about converting it to world frame for now
pusher_velocity_B = self.sample_pusher_velocity()
action_actual = pusher_velocity_B
action_zero = np.zeros(2)
episode_container = EpisodeContainer()
episode_container.set_config(env.config)
episode_container.set_name(episode_name)
obs_prev = env.get_observation()
num_timesteps = self._config['dataset'][
'num_timesteps'] + self._config['train']['n_history']
for i in xrange(num_timesteps):
if visualize:
env.render(mode='human')
# stand still for n_history timesteps
if i < self._config['train']['n_history']:
action = action_zero # just sit still
# this makes it easier to train our dynamics model
# since we need to get a history of observations
else:
exploration_type = self._config['dataset']['data_generation'][
'exploration_type']
if exploration_type == 'random':
action = self.sample_pusher_velocity()
elif exploration_type == 'constant':
action = action_actual
obs, reward, done, info = env.step(action)
episode_container.add_obs_action(obs_prev, action)
obs_prev = obs
if DEBUG:
print("\n\nsim_time = ", env._sim_time)
print("obs sim_time = ", obs['sim_time'])
return episode_container
def collect_episodes(config,
output_dir=None,
visualize=True,
debug=False,
run_from_thread=False,
seed=None):
# gets a random seed for each thread/process independently
if seed is None:
seed = np.random.RandomState().randint(0, 10000)
set_seed(seed)
if output_dir is None:
output_dir = os.path.join(os.getcwd(), 'data')
if not os.path.exists(output_dir):
os.makedirs(output_dir)
# save the config
config_save_file = os.path.join(output_dir, 'config.yaml')
save_yaml(config, config_save_file)
# initialize config for DataCollector
num_episodes = config['dataset']['num_episodes']
# record some metadata
metadata = dict()
metadata['episodes'] = dict()
while (len(metadata['episodes']) < num_episodes):
i = len(metadata['episodes'])
if debug:
input("Press Enter to continue...")
print("\n")
start_time = time.time()
print("collecting episode %d of %d" % (i + 1, num_episodes))
name = "%s_idx_%d" % (get_current_YYYY_MM_DD_hh_mm_ss_ms(), i)
n_his = config['train_dynamics']['n_history']
ic = generate_initial_condition(
config=config,
T_aug_enabled=True,
n_his=n_his,
randomize_velocity=True,
randomize_sdf=True,
randomize_color=True,
)
env = DrakeMugsEnv(ic['config'], visualize=visualize)
if debug:
print("initial condition\n", ic)
# set initial condition on environment
if visualize:
print("setting target realtime rate 1.0")
env.simulator.set_target_realtime_rate(1.0)
env.reset()
context = env.get_mutable_context()
env.set_object_position(context, ic['q_slider'])
env.set_pusher_position(context, ic['q_pusher'])
print("ic['action_sequence'].shape", ic['action_sequence'].shape)
# simulate for 10 seconds to let the mug stabilize
action_zero = env.get_zero_action()
env.step(action_zero, dt=10.0)
episode = collect_single_episode(
env, action_seq=ic['action_sequence'])['episode_container']
# potentially discard it if the object didn't move during the data collection
if len(episode._data['trajectory']) < 10:
print("trajectory was too short, skipping")
continue
obs_start = episode._data['trajectory'][0]['observation']
obs_end = episode._data['trajectory'][-1]['observation']
q_slider_start = obs_start['slider']['position']['translation']
q_slider_end = obs_end['slider']['position']['translation']
dq_slider = obs_start['slider']['position']['translation'] - obs_end[
'slider']['position']['translation']
if debug:
print("len(episode._data['trajectory'])",
len(episode._data['trajectory']))
print("q_slider_start", q_slider_start)
print("q_slider_end", q_slider_end)
print("dq_slider", dq_slider)
print("np.linalg.norm(dq_slider)", np.linalg.norm(dq_slider))
pose_error = compute_pose_error(obs_start, obs_end)
# if slider didn't move by at least 1 mm then discard this episode
if (pose_error['position_error'] <
0.01) and (pose_error['angle_error_degrees'] < 10):
print(
"discarding episode since slider didn't move sufficiently far")
continue
print("saving to disk")
metadata['episodes'][name] = dict()
image_data_file = episode.save_images_to_hdf5(output_dir)
non_image_data_file = episode.save_non_image_data_to_pickle(output_dir)
print("output_dir:", output_dir)
print("non_image_data.keys()", episode.non_image_data.keys())
metadata['episodes'][name]['non_image_data_file'] = non_image_data_file
metadata['episodes'][name]['image_data_file'] = image_data_file
print("done saving to disk")
elapsed = time.time() - start_time
print("single episode took: %.2f seconds" % (elapsed))
if not run_from_thread:
save_yaml(metadata, os.path.join(output_dir, 'metadata.yaml'))
print("Finished collecting episodes")
return {'metadata': metadata}
