def __init__(self, env_spec, embedding, task_space):
#Parameterized.__init__(self)
self._env_spec = env_spec
self._embedding = embedding
self._task_space = task_space
self._task_observation_space = concat_spaces(
self._task_space, self._env_spec.observation_space)
def run_task(v):
v = SimpleNamespace(**v)
task_names = sorted(v.tasks.keys())
task_args = [v.tasks[t]['args'] for t in task_names]
task_kwargs = [v.tasks[t]['kwargs'] for t in task_names]
# Environment
env = TfEnv(
normalize(
MultiTaskEnv(
task_env_cls=Point3dEnv,
task_args=task_args,
task_kwargs=task_kwargs)))
# Latent space and embedding specs
# TODO(gh/10): this should probably be done in Embedding or Algo
latent_lb = np.zeros(v.latent_length, )
latent_ub = np.ones(v.latent_length, )
latent_space = Box(latent_lb, latent_ub)
# trajectory space is (TRAJ_ENC_WINDOW, act_obs) where act_obs is a stacked
# vector of flattened actions and observations
act_lb, act_ub = env.action_space.bounds
act_lb_flat = env.action_space.flatten(act_lb)
act_ub_flat = env.action_space.flatten(act_ub)
obs_lb, obs_ub = env.observation_space.bounds
obs_lb_flat = env.observation_space.flatten(obs_lb)
obs_ub_flat = env.observation_space.flatten(obs_ub)
# act_obs_lb = np.concatenate([act_lb_flat, obs_lb_flat])
# act_obs_ub = np.concatenate([act_ub_flat, obs_ub_flat])
act_obs_lb = obs_lb_flat
act_obs_ub = obs_ub_flat
# act_obs_lb = act_lb_flat
# act_obs_ub = act_ub_flat
traj_lb = np.stack([act_obs_lb] * v.inference_window)
traj_ub = np.stack([act_obs_ub] * v.inference_window)
traj_space = Box(traj_lb, traj_ub)
task_embed_spec = EmbeddingSpec(env.task_space, latent_space)
traj_embed_spec = EmbeddingSpec(traj_space, latent_space)
task_obs_space = concat_spaces(env.task_space, env.observation_space)
env_spec_embed = EnvSpec(task_obs_space, env.action_space)
# TODO(): rename to inference_network
traj_embedding = GaussianMLPEmbedding(
name="inference",
embedding_spec=traj_embed_spec,
hidden_sizes=(20, 20),
std_share_network=True,
init_std=1.0,
)
# Embeddings
task_embedding = GaussianMLPEmbedding(
name="embedding",
embedding_spec=task_embed_spec,
hidden_sizes=(20, 20),
std_share_network=True,
init_std=1.0,
max_std=2.0,
# normalize=True,
)
# Multitask policy
policy = GaussianMLPMultitaskPolicy(
name="policy",
env_spec=env.spec,
task_space=env.task_space,
embedding=task_embedding,
hidden_sizes=(20, 10),
std_share_network=True,
# max_std=6.0,
init_std=6.0,
)
# baseline = MultiTaskLinearFeatureBaseline(env_spec=env_spec_embed)
extra = v.latent_length + len(v.tasks)
baseline = MultiTaskGaussianMLPBaseline(
env_spec=env.spec, extra_dims=extra)
algo = PPOTaskEmbedding(
env=env,
policy=policy,
baseline=baseline,
inference=traj_embedding,
batch_size=v.batch_size, # 4096
max_path_length=50,
n_itr=500,
discount=0.99,
step_size=0.2,
plot=True,
policy_ent_coeff=v.policy_ent_coeff,
embedding_ent_coeff=v.embedding_ent_coeff,
inference_ce_coeff=v.inference_ce_coeff,
num_tasks_held_out=1,
)
algo.train()
def run_task(v):
v = SimpleNamespace(**v)
task_names = sorted(v.tasks.keys())
task_args = [v.tasks[t]['args'] for t in task_names]
task_kwargs = [v.tasks[t]['kwargs'] for t in task_names]
# Environment
env = TfEnv(
MultiTaskEnv(task_env_cls=PointEnv,
task_args=task_args,
task_kwargs=task_kwargs))
# Latent space and embedding specs
# TODO(gh/10): this should probably be done in Embedding or Algo
latent_lb = np.zeros(v.latent_length, )
latent_ub = np.ones(v.latent_length, )
latent_space = Box(latent_lb, latent_ub)
# trajectory space is (TRAJ_ENC_WINDOW, act_obs) where act_obs is a stacked
# vector of flattened actions and observations
act_lb, act_ub = env.action_space.bounds
act_lb_flat = env.action_space.flatten(act_lb)
act_ub_flat = env.action_space.flatten(act_ub)
obs_lb, obs_ub = env.observation_space.bounds
obs_lb_flat = env.observation_space.flatten(obs_lb)
obs_ub_flat = env.observation_space.flatten(obs_ub)
# act_obs_lb = np.concatenate([act_lb_flat, obs_lb_flat])
# act_obs_ub = np.concatenate([act_ub_flat, obs_ub_flat])
act_obs_lb = obs_lb_flat
act_obs_ub = obs_ub_flat
# act_obs_lb = act_lb_flat
# act_obs_ub = act_ub_flat
traj_lb = np.stack([act_obs_lb] * v.inference_window)
traj_ub = np.stack([act_obs_ub] * v.inference_window)
traj_space = Box(traj_lb, traj_ub)
task_embed_spec = EmbeddingSpec(env.task_space, latent_space)
traj_embed_spec = EmbeddingSpec(traj_space, latent_space)
task_obs_space = concat_spaces(env.task_space, env.observation_space)
env_spec_embed = EnvSpec(task_obs_space, env.action_space)
# TODO(): rename to inference_network
traj_embedding = GaussianMLPEmbedding(
name="inference",
embedding_spec=traj_embed_spec,
hidden_sizes=(20, 10), # was the same size as policy in Karol's paper
std_share_network=True,
init_std=2.0,
mean_output_nonlinearity=tf.nn.tanh,
min_std=v.embedding_min_std,
)
# Embeddings
task_embedding = GaussianMLPEmbedding(
name="embedding",
embedding_spec=task_embed_spec,
hidden_sizes=(20, 20),
std_share_network=True,
init_std=v.embedding_init_std,
max_std=v.embedding_max_std,
mean_output_nonlinearity=tf.nn.tanh,
min_std=v.embedding_min_std,
)
# Multitask policy
policy = GaussianMLPMultitaskPolicy(
name="policy",
env_spec=env.spec,
task_space=env.task_space,
embedding=task_embedding,
hidden_sizes=(32, 16),
std_share_network=True,
max_std=v.policy_max_std,
init_std=v.policy_init_std,
min_std=v.policy_min_std,
)
extra = v.latent_length + len(v.tasks)
baseline = MultiTaskGaussianMLPBaseline(env_spec=env.spec,
extra_dims=extra)
algo = PPOTaskEmbedding(
env=env,
policy=policy,
baseline=baseline,
inference=traj_embedding,
batch_size=v.batch_size,
max_path_length=v.max_path_length,
n_itr=600,
discount=0.99,
step_size=0.2,
plot=True,
policy_ent_coeff=v.policy_ent_coeff,
embedding_ent_coeff=v.embedding_ent_coeff,
inference_ce_coeff=v.inference_ce_coeff,
use_softplus_entropy=True,
stop_ce_gradient=True,
)
algo.train()
def run_task(*_):
# Environment
env = TfEnv(
normalize(
MultiTaskEnv(task_env_cls=PR2ArmClockEnv,
task_args=TASK_ARGS,
task_kwargs=TASK_KWARGS)))
# Latent space and embedding specs
# TODO(gh/10): this should probably be done in Embedding or Algo
latent_lb = np.zeros(LATENT_LENGTH, )
latent_ub = np.ones(LATENT_LENGTH, )
latent_space = Box(latent_lb, latent_ub)
# trajectory space is (TRAJ_ENC_WINDOW, act_obs) where act_obs is a stacked
# vector of flattened actions and observations
act_lb, act_ub = env.action_space.bounds
act_lb_flat = env.action_space.flatten(act_lb)
act_ub_flat = env.action_space.flatten(act_ub)
obs_lb, obs_ub = env.observation_space.bounds
obs_lb_flat = env.observation_space.flatten(obs_lb)
obs_ub_flat = env.observation_space.flatten(obs_ub)
act_obs_lb = np.concatenate([act_lb_flat, obs_lb_flat])
act_obs_ub = np.concatenate([act_ub_flat, obs_ub_flat])
traj_lb = np.stack([act_obs_lb] * TRAJ_ENC_WINDOW)
traj_ub = np.stack([act_obs_ub] * TRAJ_ENC_WINDOW)
traj_space = Box(traj_lb, traj_ub)
task_embed_spec = EmbeddingSpec(env.task_space, latent_space)
traj_embed_spec = EmbeddingSpec(traj_space, latent_space)
latent_obs_space = concat_spaces(latent_space, env.observation_space)
env_spec_embed = EnvSpec(latent_obs_space, env.action_space)
# Base policy
policy = GaussianMLPPolicy(
name="policy",
env_spec=env_spec_embed,
hidden_sizes=(32, 32),
adaptive_std=True, # Must be True for embedding learning
)
# Embeddings
task_embedding = GaussianMLPEmbedding(
name="task_embedding",
embedding_spec=task_embed_spec,
hidden_sizes=(32, 32),
adaptive_std=True, # Must be True for embedding learning
)
traj_embedding = GaussianMLPEmbedding(
name="traj_embedding",
embedding_spec=traj_embed_spec,
hidden_sizes=(32, 32),
adaptive_std=True, # Must be True for embedding learning
)
baseline = LinearFeatureBaseline(env_spec=env_spec_embed)
algo = TRPOTaskEmbedding(
env=env,
policy=policy,
baseline=baseline,
embedding=task_embedding,
inference=traj_embedding,
batch_size=4000,
max_path_length=MAX_PATH_LENGTH,
n_itr=400000000,
discount=0.99,
step_size=0.01,
plot=True,
)
algo.train()
def run_task(plot=True, *_):
set_seed(0)
# Environment (train on light point mass)
from sandbox.embed2learn.envs import point_mass_env
from dm_control import suite
suite._DOMAINS["embed_point_mass"] = point_mass_env
env = TfEnv(
normalize(
MultiTaskEnv(
task_env_cls=DmControlEnv,
task_args=[["embed_point_mass", "light"]] * len(TASK_NAMES),
task_kwargs=TASK_KWARGS)))
# Latent space and embedding specs
# TODO(gh/10): this should probably be done in Embedding or Algo
latent_lb = np.zeros(LATENT_LENGTH, )
latent_ub = np.ones(LATENT_LENGTH, )
latent_space = Box(latent_lb, latent_ub)
# trajectory space is (TRAJ_ENC_WINDOW, act_obs) where act_obs is a stacked
# vector of flattened actions and observations
act_lb, act_ub = env.action_space.bounds
act_lb_flat = env.action_space.flatten(act_lb)
act_ub_flat = env.action_space.flatten(act_ub)
obs_lb, obs_ub = env.observation_space.bounds
obs_lb_flat = env.observation_space.flatten(obs_lb)
obs_ub_flat = env.observation_space.flatten(obs_ub)
#act_obs_lb = np.concatenate([act_lb_flat, obs_lb_flat])
#act_obs_ub = np.concatenate([act_ub_flat, obs_ub_flat])
act_obs_lb = obs_lb_flat
act_obs_ub = obs_ub_flat
traj_lb = np.stack([act_obs_lb] * TRAJ_ENC_WINDOW)
traj_ub = np.stack([act_obs_ub] * TRAJ_ENC_WINDOW)
traj_space = Box(traj_lb, traj_ub)
task_embed_spec = EmbeddingSpec(env.task_space, latent_space)
traj_embed_spec = EmbeddingSpec(traj_space, latent_space)
task_obs_space = concat_spaces(env.task_space, env.observation_space)
env_spec_embed = EnvSpec(task_obs_space, env.action_space)
# Embeddings
task_embedding = GaussianMLPEmbedding(
name="task_embedding",
embedding_spec=task_embed_spec,
hidden_sizes=(20, 20),
adaptive_std=True,
init_std=0.5, # TODO was 100
max_std=0.6, # TODO find appropriate value
)
# TODO(): rename to inference_network
traj_embedding = GaussianMLPEmbedding(
name="traj_embedding",
embedding_spec=traj_embed_spec,
hidden_sizes=(20, 10), # was the same size as policy in Karol's paper
# adaptive_std=True, # Must be True for embedding learning
std_share_network=True,
init_std=0.001,
)
# Multitask policy
policy = GaussianMLPMultitaskPolicy(
name="policy",
env_spec=env.spec,
task_space=env.task_space,
embedding=task_embedding,
hidden_sizes=(20, 10),
adaptive_std=True, # Must be True for embedding learning
init_std=0.5, # TODO was 100
)
baseline = LinearFeatureBaseline(env_spec=env_spec_embed)
algo = TRPOTaskEmbedding(
env=env,
policy=policy,
baseline=baseline,
inference=traj_embedding,
batch_size=4000,
max_path_length=100,
n_itr=500,
discount=0.99,
step_size=0.01,
plot=plot,
plot_warmup_itrs=20,
policy_ent_coeff=0.0, # 0.001, #0.1,
embedding_ent_coeff=0.0, #0.1,
inference_ce_ent_coeff=0., # 0.03, #0.1, # 0.1,
)
algo.train()
def run_task(*_):
set_seed(1)
# Environment
env = TfEnv(
MultiTaskEnv(task_env_cls=PointEnv,
task_args=TASK_ARGS,
task_kwargs=TASK_KWARGS))
# Latent space and embedding specs
# TODO(gh/10): this should probably be done in Embedding or Algo
latent_lb = np.zeros(LATENT_LENGTH, )
latent_ub = np.ones(LATENT_LENGTH, )
latent_space = Box(latent_lb, latent_ub)
# trajectory space is (TRAJ_ENC_WINDOW, act_obs) where act_obs is a stacked
# vector of flattened actions and observations
act_lb, act_ub = env.action_space.bounds
act_lb_flat = env.action_space.flatten(act_lb)
act_ub_flat = env.action_space.flatten(act_ub)
obs_lb, obs_ub = env.observation_space.bounds
obs_lb_flat = env.observation_space.flatten(obs_lb)
obs_ub_flat = env.observation_space.flatten(obs_ub)
# act_obs_lb = np.concatenate([act_lb_flat, obs_lb_flat])
# act_obs_ub = np.concatenate([act_ub_flat, obs_ub_flat])
act_obs_lb = obs_lb_flat
act_obs_ub = obs_ub_flat
# act_obs_lb = act_lb_flat
# act_obs_ub = act_ub_flat
traj_lb = np.stack([act_obs_lb] * TRAJ_ENC_WINDOW)
traj_ub = np.stack([act_obs_ub] * TRAJ_ENC_WINDOW)
traj_space = Box(traj_lb, traj_ub)
task_embed_spec = EmbeddingSpec(env.task_space, latent_space)
traj_embed_spec = EmbeddingSpec(traj_space, latent_space)
task_obs_space = concat_spaces(env.task_space, env.observation_space)
env_spec_embed = EnvSpec(task_obs_space, env.action_space)
# Embeddings
task_embedding = GaussianMLPEmbedding(
name="embedding",
embedding_spec=task_embed_spec,
hidden_sizes=(20, 20),
std_share_network=True,
init_std=3.0, # 2.0
)
# TODO(): rename to inference_network
traj_embedding = GaussianMLPEmbedding(
name="inference",
embedding_spec=traj_embed_spec,
hidden_sizes=(20, 10), # was the same size as policy in Karol's paper
std_share_network=True,
)
# Multitask policy
policy = GaussianMLPMultitaskPolicy(
name="policy",
env_spec=env.spec,
task_space=env.task_space,
embedding=task_embedding,
hidden_sizes=(20, 10),
std_share_network=True, # Must be True for embedding learning
init_std=6.0, # 4.5 6.0
)
baseline = MultiTaskLinearFeatureBaseline(env_spec=env_spec_embed)
algo = TRPOTaskEmbedding(
env=env,
policy=policy,
baseline=baseline,
inference=traj_embedding,
batch_size=20000,
max_path_length=50,
n_itr=1000,
discount=0.99,
step_size=0.2,
plot=False,
policy_ent_coeff=1e-7, # 1e-7
embedding_ent_coeff=1e-3, # 1e-3
inference_ce_coeff=1e-7, # 1e-7
# kl_constraint=KLConstraint.SOFT,
# optimizer_args=dict(max_penalty=1e9),
)
algo.train()