def generate_vae_dataset(
N=10000, test_p=0.9, use_cached=True, imsize=84, show=False,
dataset_path=None, env_class=None, env_kwargs=None, init_camera=sawyer_door_env_camera,
):
filename = "/tmp/sawyer_door_push_open_and_reach" + str(N) + ".npy"
info = {}
if dataset_path is not None:
filename = local_path_from_s3_or_local_path(dataset_path)
dataset = np.load(filename)
elif use_cached and osp.isfile(filename):
dataset = np.load(filename)
print("loaded data from saved file", filename)
else:
env = env_class(**env_kwargs)
env = ImageEnv(
env, imsize,
transpose=True,
init_camera=init_camera,
normalize=True,
)
oracle_sampled_data = int(N/2)
dataset = np.zeros((N, imsize * imsize * 3))
print('Goal Space Sampling')
for i in range(oracle_sampled_data):
goal = env.sample_goal()
env.set_to_goal(goal)
img = env._get_flat_img()
dataset[i, :] = img
if show:
cv2.imshow('img', img.reshape(3, 84, 84).transpose())
cv2.waitKey(1)
print(i)
env._wrapped_env.min_y_pos=.6
policy = RandomPolicy(env.action_space)
es = OUStrategy(action_space=env.action_space, theta=0)
exploration_policy = PolicyWrappedWithExplorationStrategy(
exploration_strategy=es,
policy=policy,
)
print('Random Sampling')
for i in range(oracle_sampled_data, N):
if i % 20==0:
env.reset()
exploration_policy.reset()
for _ in range(10):
action = exploration_policy.get_action()[0]
env.wrapped_env.step(
action
)
img = env._get_flat_img()
dataset[i, :] = img
if show:
cv2.imshow('img', img.reshape(3, 84, 84).transpose())
cv2.waitKey(1)
print(i)
n = int(N * test_p)
train_dataset = dataset[:n, :]
test_dataset = dataset[n:, :]
return train_dataset, test_dataset, info
def generate_vae_dataset(
N=10000,
test_p=0.9,
use_cached=False,
imsize=84,
show=False,
dataset_path=None,
env_class=SawyerReachTorqueEnv,
env_kwargs=None,
init_camera=sawyer_torque_reacher_camera,
):
filename = "/tmp/sawyer_torque_data" + str(N) + ".npy"
info = {}
if dataset_path is not None:
filename = local_path_from_s3_or_local_path(dataset_path)
dataset = np.load(filename)
elif use_cached and osp.isfile(filename):
dataset = np.load(filename)
print("loaded data from saved file", filename)
else:
now = time.time()
if env_kwargs == None:
env_kwargs = dict()
env = env_class(**env_kwargs)
env = ImageEnv(
env,
imsize,
transpose=True,
init_camera=init_camera,
normalize=True,
)
info['env'] = env
policy = RandomPolicy(env.action_space)
es = OUStrategy(action_space=env.action_space, theta=0)
exploration_policy = PolicyWrappedWithExplorationStrategy(
exploration_strategy=es,
policy=policy,
)
dataset = np.zeros((N, imsize * imsize * 3), dtype=np.uint8)
for i in range(N):
if i % 50 == 0:
print('Reset')
env.reset_model()
exploration_policy.reset()
for _ in range(1):
action = exploration_policy.get_action()[0] * 1 / 10
env.wrapped_env.step(action)
img = env._get_flat_img()
dataset[i, :] = unormalize_image(img)
if show:
cv2.imshow('img', img.reshape(3, 84, 84).transpose())
cv2.waitKey(1)
print(i)
print("done making training data", time.time() - now)
np.save(filename, dataset)
n = int(N * test_p)
train_dataset = dataset[:n, :]
test_dataset = dataset[n:, :]
return train_dataset, test_dataset, info
def generate_uniform_dataset_reacher(env_class=None,
env_kwargs=None,
num_imgs=1000,
use_cached_dataset=False,
init_camera=None,
imsize=48,
show=False,
save_file_prefix=None,
env_id=None,
tag='',
dataset_path=None):
if dataset_path is not None:
dataset = load_local_or_remote_file(dataset_path)
return dataset
import gym
from gym.envs import registration
# trigger registration
import multiworld.envs.pygame
import multiworld.envs.mujoco
if not env_class or not env_kwargs:
env = gym.make(env_id)
else:
env = env_class(**env_kwargs)
env = ImageEnv(
env,
imsize,
init_camera=init_camera,
transpose=True,
normalize=True,
)
env.non_presampled_goal_img_is_garbage = True
if save_file_prefix is None and env_id is not None:
save_file_prefix = env_id
filename = "/tmp/{}_N{}_imsize{}uniform_images_{}.npy".format(
save_file_prefix,
str(num_imgs),
env.imsize,
tag,
)
if use_cached_dataset and osp.isfile(filename):
images = np.load(filename)
print("Loaded data from {}".format(filename))
return images
print('Sampling Uniform Dataset')
dataset = np.zeros((num_imgs, 3 * env.imsize**2), dtype=np.uint8)
for j in range(num_imgs):
obs = env.reset()
env.set_to_goal(env.get_goal())
img_f = env._get_flat_img()
if show:
img = img_f.reshape(3, env.imsize, env.imsize).transpose()
img = img[::-1, :, ::-1]
cv2.imshow('img', img)
cv2.waitKey(1)
print(j)
dataset[j, :] = unormalize_image(img_f)
np.save(filename, dataset)
print("Saving file to {}".format(filename))
return dataset
def generate_vae_dataset(
N=10000, test_p=0.9, use_cached=True, imsize=84, show=False,
dataset_path=None, policy_path=None, action_space_sampling=False, env_class=SawyerDoorEnv, env_kwargs=None,
init_camera=sawyer_door_env_camera_v2,
):
if policy_path is not None:
filename = "/tmp/sawyer_door_pull_open_oracle+random_policy_data_closer_zoom_action_limited" + str(N) + ".npy"
elif action_space_sampling:
filename = "/tmp/sawyer_door_pull_open_zoomed_in_action_space_sampling" + str(N) + ".npy"
else:
filename = "/tmp/sawyer_door_pull_open" + str(N) + ".npy"
info = {}
if dataset_path is not None:
filename = local_path_from_s3_or_local_path(dataset_path)
dataset = np.load(filename)
elif use_cached and osp.isfile(filename):
dataset = np.load(filename)
print("loaded data from saved file", filename)
else:
now = time.time()
env = env_class(**env_kwargs)
env = ImageEnv(
env, imsize,
transpose=True,
init_camera=init_camera,
normalize=True,
)
info['env'] = env
policy = RandomPolicy(env.action_space)
es = OUStrategy(action_space=env.action_space, theta=0)
exploration_policy = PolicyWrappedWithExplorationStrategy(
exploration_strategy=es,
policy=policy,
)
env.wrapped_env.reset()
dataset = np.zeros((N, imsize * imsize * 3), dtype=np.uint8)
for i in range(N):
if i % 20==0:
env.reset_model()
exploration_policy.reset()
for _ in range(10):
action = exploration_policy.get_action()[0]
env.wrapped_env.step(
action
)
# env.set_to_goal_angle(env.get_goal()['state_desired_goal'])
img = env._get_flat_img()
dataset[i, :] = unormalize_image(img)
if show:
cv2.imshow('img', img.reshape(3, 84, 84).transpose())
cv2.waitKey(1)
print(i)
print("done making training data", filename, time.time() - now)
np.save(filename, dataset)
n = int(N * test_p)
train_dataset = dataset[:n, :]
test_dataset = dataset[n:, :]
return train_dataset, test_dataset, info
def generate_vae_dataset(
N=10000,
test_p=0.9,
use_cached=True,
imsize=84,
show=False,
dataset_path=None,
policy_path=None,
action_space_sampling=False,
env_class=SawyerPushAndPullDoorEnv,
env_kwargs=None,
action_plus_random_sampling=False,
init_camera=sawyer_door_env_camera,
ratio_action_sample_to_random=1 / 2,
env_id=None,
):
if policy_path is not None:
filename = "/tmp/sawyer_door_push_and_pull_open_oracle+random_policy_data_closer_zoom_action_limited" + str(
N) + ".npy"
elif action_space_sampling:
filename = "/tmp/sawyer_door_push_and_pull_open_zoomed_in_action_space_sampling" + str(
N) + ".npy"
else:
filename = "/tmp/sawyer_door_push_and_pull_open" + str(N) + ".npy"
info = {}
if dataset_path is not None:
filename = local_path_from_s3_or_local_path(dataset_path)
dataset = np.load(filename)
elif use_cached and osp.isfile(filename):
dataset = np.load(filename)
print("loaded data from saved file", filename)
elif action_plus_random_sampling:
if env_id is not None:
import gym
env = gym.make(env_id)
else:
env = env_class(**env_kwargs)
env = ImageEnv(
env,
imsize,
transpose=True,
init_camera=init_camera,
normalize=True,
)
action_sampled_data = int(N * ratio_action_sample_to_random)
dataset = np.zeros((N, imsize * imsize * 3), dtype=np.uint8)
print('Action Space Sampling')
for i in range(action_sampled_data):
goal = env.sample_goal()
env.set_to_goal(goal)
img = env._get_flat_img()
dataset[i, :] = unormalize_image(img)
if show:
cv2.imshow('img', img.reshape(3, 84, 84).transpose())
cv2.waitKey(1)
print(i)
policy = RandomPolicy(env.action_space)
es = OUStrategy(action_space=env.action_space, theta=0)
exploration_policy = PolicyWrappedWithExplorationStrategy(
exploration_strategy=es,
policy=policy,
)
print('Random Sampling')
for i in range(action_sampled_data, N):
if i % 20 == 0:
env.reset()
exploration_policy.reset()
for _ in range(10):
action = exploration_policy.get_action()[0]
env.wrapped_env.step(action)
goal = env.sample_goal()
env.set_to_goal_angle(goal['state_desired_goal'])
img = env._get_flat_img()
dataset[i, :] = unormalize_image(img)
if show:
cv2.imshow('img', img.reshape(3, 84, 84).transpose())
cv2.waitKey(1)
print(i)
env._wrapped_env.min_y_pos = .5
info['env'] = env
else:
raise NotImplementedError()
n = int(N * test_p)
train_dataset = dataset[:n, :]
test_dataset = dataset[n:, :]
return train_dataset, test_dataset, info
def generate_vae_dataset(
N=10000,
test_p=0.9,
use_cached=True,
imsize=84,
show=False,
dataset_path=None,
action_space_sampling=False,
init_camera=None,
env_class=None,
env_kwargs=None,
):
filename = "/tmp/sawyer_xyz_pos_control_new_zoom_cam" + str(N) + '.npy'
info = {}
if dataset_path is not None:
filename = local_path_from_s3_or_local_path(dataset_path)
dataset = np.load(filename)
elif use_cached and osp.isfile(filename):
dataset = np.load(filename)
print("loaded data from saved file", filename)
else:
now = time.time()
if env_kwargs == None:
env_kwargs = dict()
env = env_class(**env_kwargs)
env = ImageEnv(
env,
imsize,
transpose=True,
init_camera=init_camera,
normalize=True,
)
dataset = np.zeros((N, imsize * imsize * 3), dtype=np.uint8)
if action_space_sampling:
action_space = Box(np.array([-.1, .5, 0]), np.array([.1, .7, .5]))
for i in range(N):
env.set_to_goal(env.sample_goal())
img = env._get_flat_img()
dataset[i, :] = unormalize_image(img)
if show:
cv2.imshow('img', img.reshape(3, 84, 84).transpose())
cv2.waitKey(1)
print(i)
info['env'] = env
else:
policy = RandomPolicy(env.action_space)
es = OUStrategy(action_space=env.action_space, theta=0)
exploration_policy = PolicyWrappedWithExplorationStrategy(
exploration_strategy=es,
policy=policy,
)
for i in range(N):
# Move the goal out of the image
env.wrapped_env.set_goal(np.array([100, 100, 100]))
if i % 50 == 0:
print('Reset')
env.reset()
exploration_policy.reset()
for _ in range(1):
action = exploration_policy.get_action()[0] * 10
env.wrapped_env.step(action)
img = env.step(env.action_space.sample())[0]
dataset[i, :] = img
if show:
cv2.imshow('img', img.reshape(3, 84, 84).transpose())
cv2.waitKey(1)
print(i)
print("done making training data", time.time() - now)
np.save(filename, dataset)
n = int(N * test_p)
train_dataset = dataset[:n, :]
test_dataset = dataset[n:, :]
return train_dataset, test_dataset, info
