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(variant):
print(variant)
from tqdm import tqdm
env_class = variant.get('env_class', None)
env_kwargs = variant.get('env_kwargs', None)
env_id = variant.get('env_id', None)
N = variant.get('N', 10000)
batch_size = variant.get('batch_size', 128)
test_p = variant.get('test_p', 0.9)
use_cached = variant.get('use_cached', True)
imsize = variant.get('imsize', 84)
num_channels = variant.get('num_channels', 3)
show = variant.get('show', False)
init_camera = variant.get('init_camera', None)
dataset_path = variant.get('dataset_path', None)
augment_data = variant.get('augment_data', False)
data_filter_fn = variant.get('data_filter_fn', lambda x: x)
delete_after_loading = variant.get('delete_after_loading', False)
oracle_dataset_using_set_to_goal = variant.get(
'oracle_dataset_using_set_to_goal', False)
random_rollout_data = variant.get('random_rollout_data', False)
random_rollout_data_set_to_goal = variant.get(
'random_rollout_data_set_to_goal', True)
random_and_oracle_policy_data = variant.get(
'random_and_oracle_policy_data', False)
random_and_oracle_policy_data_split = variant.get(
'random_and_oracle_policy_data_split', 0)
policy_file = variant.get('policy_file', None)
n_random_steps = variant.get('n_random_steps', 100)
vae_dataset_specific_env_kwargs = variant.get(
'vae_dataset_specific_env_kwargs', None)
save_file_prefix = variant.get('save_file_prefix', None)
non_presampled_goal_img_is_garbage = variant.get(
'non_presampled_goal_img_is_garbage', None)
conditional_vae_dataset = variant.get('conditional_vae_dataset', False)
use_env_labels = variant.get('use_env_labels', False)
use_linear_dynamics = variant.get('use_linear_dynamics', False)
enviorment_dataset = variant.get('enviorment_dataset', False)
save_trajectories = variant.get('save_trajectories', False)
save_trajectories = save_trajectories or use_linear_dynamics or conditional_vae_dataset
tag = variant.get('tag', '')
assert N % n_random_steps == 0, "Fix N/horizon or dataset generation will fail"
from multiworld.core.image_env import ImageEnv, unormalize_image
import rlkit.torch.pytorch_util as ptu
from rlkit.util.io import load_local_or_remote_file
from rlkit.data_management.dataset import (
TrajectoryDataset, ImageObservationDataset, InitialObservationDataset,
EnvironmentDataset, ConditionalDynamicsDataset,
InitialObservationNumpyDataset, InfiniteBatchLoader,
InitialObservationNumpyJitteringDataset)
info = {}
use_test_dataset = False
if dataset_path is not None:
if type(dataset_path) == str:
dataset = load_local_or_remote_file(
dataset_path, delete_after_loading=delete_after_loading)
dataset = dataset.item()
N = dataset['observations'].shape[0] * dataset[
'observations'].shape[1]
n_random_steps = dataset['observations'].shape[1]
if isinstance(dataset_path, list):
dataset = concatenate_datasets(dataset_path)
N = dataset['observations'].shape[0] * dataset[
'observations'].shape[1]
n_random_steps = dataset['observations'].shape[1]
if isinstance(dataset_path, dict):
if type(dataset_path['train']) == str:
dataset = load_local_or_remote_file(
dataset_path['train'],
delete_after_loading=delete_after_loading)
dataset = dataset.item()
elif isinstance(dataset_path['train'], list):
dataset = concatenate_datasets(dataset_path['train'])
if type(dataset_path['test']) == str:
test_dataset = load_local_or_remote_file(
dataset_path['test'],
delete_after_loading=delete_after_loading)
test_dataset = test_dataset.item()
elif isinstance(dataset_path['test'], list):
test_dataset = concatenate_datasets(dataset_path['test'])
N = dataset['observations'].shape[0] * dataset[
'observations'].shape[1]
n_random_steps = dataset['observations'].shape[1]
use_test_dataset = True
else:
if env_kwargs is None:
env_kwargs = {}
if save_file_prefix is None:
save_file_prefix = env_id
if save_file_prefix is None:
save_file_prefix = env_class.__name__
filename = "/tmp/{}_N{}_{}_imsize{}_random_oracle_split_{}{}.npy".format(
save_file_prefix,
str(N),
init_camera.__name__
if init_camera and hasattr(init_camera, '__name__') else '',
imsize,
random_and_oracle_policy_data_split,
tag,
)
if use_cached and osp.isfile(filename):
dataset = load_local_or_remote_file(
filename, delete_after_loading=delete_after_loading)
if conditional_vae_dataset:
dataset = dataset.item()
print("loaded data from saved file", filename)
else:
now = time.time()
if env_id is not None:
import gym
import multiworld
multiworld.register_all_envs()
env = gym.make(env_id)
else:
if vae_dataset_specific_env_kwargs is None:
vae_dataset_specific_env_kwargs = {}
for key, val in env_kwargs.items():
if key not in vae_dataset_specific_env_kwargs:
vae_dataset_specific_env_kwargs[key] = val
env = env_class(**vae_dataset_specific_env_kwargs)
if not isinstance(env, ImageEnv):
env = ImageEnv(
env,
imsize,
init_camera=init_camera,
transpose=True,
normalize=True,
non_presampled_goal_img_is_garbage=
non_presampled_goal_img_is_garbage,
)
else:
imsize = env.imsize
env.non_presampled_goal_img_is_garbage = non_presampled_goal_img_is_garbage
env.reset()
info['env'] = env
if random_and_oracle_policy_data:
policy_file = load_local_or_remote_file(policy_file)
policy = policy_file['policy']
policy.to(ptu.device)
if random_rollout_data:
from rlkit.exploration_strategies.ou_strategy import OUStrategy
policy = OUStrategy(env.action_space)
if save_trajectories:
dataset = {
'observations':
np.zeros((N // n_random_steps, n_random_steps,
imsize * imsize * num_channels),
dtype=np.uint8),
'actions':
np.zeros((N // n_random_steps, n_random_steps,
env.action_space.shape[0]),
dtype=np.float),
'env':
np.zeros(
(N // n_random_steps, imsize * imsize * num_channels),
dtype=np.uint8),
}
else:
dataset = np.zeros((N, imsize * imsize * num_channels),
dtype=np.uint8)
labels = []
for i in tqdm(range(N)):
if random_and_oracle_policy_data:
num_random_steps = int(N *
random_and_oracle_policy_data_split)
if i < num_random_steps:
env.reset()
for _ in range(n_random_steps):
obs = env.step(env.action_space.sample())[0]
else:
obs = env.reset()
policy.reset()
for _ in range(n_random_steps):
policy_obs = np.hstack((
obs['state_observation'],
obs['state_desired_goal'],
))
action, _ = policy.get_action(policy_obs)
obs, _, _, _ = env.step(action)
elif random_rollout_data: #ADD DATA WHERE JUST PUCK MOVES
if i % n_random_steps == 0:
env.reset()
policy.reset()
env_img = env._get_obs()['image_observation']
if random_rollout_data_set_to_goal:
env.set_to_goal(env.get_goal())
obs = env._get_obs()
u = policy.get_action_from_raw_action(
env.action_space.sample())
env.step(u)
elif oracle_dataset_using_set_to_goal:
print(i)
goal = env.sample_goal()
env.set_to_goal(goal)
obs = env._get_obs()
else:
env.reset()
for _ in range(n_random_steps):
obs = env.step(env.action_space.sample())[0]
img = obs['image_observation']
if use_env_labels:
labels.append(obs['label'])
if save_trajectories:
dataset['observations'][
i // n_random_steps,
i % n_random_steps, :] = unormalize_image(img)
dataset['actions'][i // n_random_steps,
i % n_random_steps, :] = u
dataset['env'][i // n_random_steps, :] = unormalize_image(
env_img)
else:
dataset[i, :] = unormalize_image(img)
if show:
img = img.reshape(3, imsize, imsize).transpose()
img = img[::-1, :, ::-1]
cv2.imshow('img', img)
cv2.waitKey(1)
# radius = input('waiting...')
print("done making training data", filename, time.time() - now)
np.save(filename, dataset)
#np.save(filename[:-4] + 'labels.npy', np.array(labels))
info['train_labels'] = []
info['test_labels'] = []
dataset = data_filter_fn(dataset)
if use_linear_dynamics and conditional_vae_dataset:
num_trajectories = N // n_random_steps
n = int(num_trajectories * test_p)
train_dataset = ConditionalDynamicsDataset({
'observations':
dataset['observations'][:n, :, :],
'actions':
dataset['actions'][:n, :, :],
'env':
dataset['env'][:n, :]
})
test_dataset = ConditionalDynamicsDataset({
'observations':
dataset['observations'][n:, :, :],
'actions':
dataset['actions'][n:, :, :],
'env':
dataset['env'][n:, :]
})
num_trajectories = N // n_random_steps
n = int(num_trajectories * test_p)
indices = np.arange(num_trajectories)
np.random.shuffle(indices)
train_i, test_i = indices[:n], indices[n:]
try:
train_dataset = ConditionalDynamicsDataset({
'observations':
dataset['observations'][train_i, :, :],
'actions':
dataset['actions'][train_i, :, :],
'env':
dataset['env'][train_i, :]
})
test_dataset = ConditionalDynamicsDataset({
'observations':
dataset['observations'][test_i, :, :],
'actions':
dataset['actions'][test_i, :, :],
'env':
dataset['env'][test_i, :]
})
except:
train_dataset = ConditionalDynamicsDataset({
'observations':
dataset['observations'][train_i, :, :],
'actions':
dataset['actions'][train_i, :, :],
})
test_dataset = ConditionalDynamicsDataset({
'observations':
dataset['observations'][test_i, :, :],
'actions':
dataset['actions'][test_i, :, :],
})
elif use_linear_dynamics:
num_trajectories = N // n_random_steps
n = int(num_trajectories * test_p)
train_dataset = TrajectoryDataset({
'observations':
dataset['observations'][:n, :, :],
'actions':
dataset['actions'][:n, :, :]
})
test_dataset = TrajectoryDataset({
'observations':
dataset['observations'][n:, :, :],
'actions':
dataset['actions'][n:, :, :]
})
elif enviorment_dataset:
n = int(n_random_steps * test_p)
train_dataset = EnvironmentDataset({
'observations':
dataset['observations'][:, :n, :],
})
test_dataset = EnvironmentDataset({
'observations':
dataset['observations'][:, n:, :],
})
elif conditional_vae_dataset:
num_trajectories = N // n_random_steps
n = int(num_trajectories * test_p)
indices = np.arange(num_trajectories)
np.random.shuffle(indices)
train_i, test_i = indices[:n], indices[n:]
if augment_data:
dataset_class = InitialObservationNumpyJitteringDataset
else:
dataset_class = InitialObservationNumpyDataset
if 'env' not in dataset:
dataset['env'] = dataset['observations'][:, 0]
if use_test_dataset and ('env' not in test_dataset):
test_dataset['env'] = test_dataset['observations'][:, 0]
if use_test_dataset:
train_dataset = dataset_class({
'observations':
dataset['observations'],
'env':
dataset['env']
})
test_dataset = dataset_class({
'observations':
test_dataset['observations'],
'env':
test_dataset['env']
})
else:
train_dataset = dataset_class({
'observations':
dataset['observations'][train_i, :, :],
'env':
dataset['env'][train_i, :]
})
test_dataset = dataset_class({
'observations':
dataset['observations'][test_i, :, :],
'env':
dataset['env'][test_i, :]
})
train_batch_loader_kwargs = variant.get(
'train_batch_loader_kwargs',
dict(
batch_size=batch_size,
num_workers=0,
))
test_batch_loader_kwargs = variant.get(
'test_batch_loader_kwargs',
dict(
batch_size=batch_size,
num_workers=0,
))
train_data_loader = data.DataLoader(train_dataset,
shuffle=True,
drop_last=True,
**train_batch_loader_kwargs)
test_data_loader = data.DataLoader(test_dataset,
shuffle=True,
drop_last=True,
**test_batch_loader_kwargs)
train_dataset = InfiniteBatchLoader(train_data_loader)
test_dataset = InfiniteBatchLoader(test_data_loader)
else:
n = int(N * test_p)
train_dataset = ImageObservationDataset(dataset[:n, :])
test_dataset = ImageObservationDataset(dataset[n:, :])
return train_dataset, test_dataset, info
