def _make_trunk(
self,
input_dim: int,
hidden_dim: int,
output_dim: int,
num_layers: int,
multitask_cfg: ConfigType,
) -> ModelType:
if ("actor_cfg" in multitask_cfg and multitask_cfg.actor_cfg
and "moe_cfg" in multitask_cfg.actor_cfg
and multitask_cfg.actor_cfg.moe_cfg.should_use):
moe_cfg = multitask_cfg.actor_cfg.moe_cfg
if moe_cfg.mode == "soft_modularization":
trunk = SoftModularizedMLP(
num_experts=moe_cfg.num_experts,
in_features=input_dim,
out_features=output_dim,
num_layers=2,
hidden_features=hidden_dim,
bias=True,
)
else:
raise NotImplementedError(
f"""`moe_cfg.mode` = {moe_cfg.mode} is not implemented.""")
else:
trunk = agent_utils.build_mlp( # type: ignore[assignment]
input_dim=input_dim,
hidden_dim=hidden_dim,
output_dim=output_dim,
num_layers=num_layers,
)
# This seems to be a false alarm since both nn.Module and
# SoftModularizedMLP are subtypes of ModelType.
return trunk
def __init__(
self,
env_obs_shape: List[int],
multitask_cfg: ConfigType,
feature_dim: int,
num_layers: int,
hidden_dim: int,
should_tie_encoders: bool,
):
"""Feedforward encoder for state observations.
Args:
env_obs_shape (List[int]): shape of the observation that the actor gets.
multitask_cfg (ConfigType): config for encoding the multitask knowledge.
feature_dim (int): feature dimension.
num_layers (int, optional): number of layers. Defaults to 2.
hidden_dim (int, optional): number of conv filters per layer. Defaults to 32.
should_tie_encoders (bool): should the feed-forward layers be tied.
"""
super().__init__(env_obs_shape=env_obs_shape,
multitask_cfg=multitask_cfg)
assert len(env_obs_shape) == 1
self.num_layers = num_layers
self.trunk = agent_utils.build_mlp(
input_dim=env_obs_shape[0],
hidden_dim=hidden_dim,
num_layers=num_layers,
output_dim=feature_dim,
)
self.should_tie_encoders = should_tie_encoders
def _make_trunk(
self,
obs_dim: int,
action_dim: int,
hidden_dim: int,
output_dim: int,
num_layers: int,
multitask_cfg: ConfigType,
) -> ModelType:
"""Make the tunk for the Q-function.
Args:
obs_dim (int): size of the observation.
action_dim (int): size of the action vector.
hidden_dim (int): size of the hidden layer of the trunk.
output_dim (int): size of the output.
num_layers (int): number of layers in the model.
multitask_cfg (ConfigType): config for encoding the multitask knowledge.
Returns:
ModelType:
"""
if (
"critic_cfg" in multitask_cfg
and multitask_cfg.critic_cfg
and "moe_cfg" in multitask_cfg.critic_cfg
and multitask_cfg.critic_cfg.moe_cfg.should_use
):
moe_cfg = multitask_cfg.critic_cfg.moe_cfg
if moe_cfg.mode == "soft_modularization":
trunk = SoftModularizedMLP(
num_experts=moe_cfg.num_experts,
in_features=obs_dim,
out_features=output_dim,
num_layers=2,
hidden_features=hidden_dim,
bias=True,
)
else:
raise NotImplementedError(
f"""`moe_cfg.mode` = {moe_cfg.mode} is not implemented."""
)
else:
trunk = agent_utils.build_mlp( # type: ignore[assignment]
input_dim=obs_dim + action_dim,
hidden_dim=hidden_dim,
output_dim=output_dim,
num_layers=num_layers,
)
# This seems to be a false alarm since both nn.Module and
# SoftModularizedMLP are subtypes of ModelType.
return trunk
def __init__(
self,
num_tasks: int,
num_experts: int,
embedding_dim: int,
hidden_dim: int,
num_layers: int,
temperature: bool,
should_use_soft_attention: bool,
task_encoder_cfg: ConfigType,
multitask_cfg: ConfigType,
topk: Optional[int] = None,
):
super().__init__(multitask_cfg=multitask_cfg)
self.should_use_soft_attention = should_use_soft_attention
self.temperature = temperature
self.should_use_task_encoding = task_encoder_cfg.should_use_task_encoding
self.should_detach_task_encoding = (
self.should_use_task_encoding
and task_encoder_cfg.should_detach_task_encoding)
if not self.should_use_task_encoding:
self.emb = nn.Embedding(
num_embeddings=num_tasks,
embedding_dim=embedding_dim,
)
self.emb.apply(agent_utils.weight_init)
else:
embedding_dim = multitask_cfg.task_encoder_cfg.model_cfg.embedding_dim
self.trunk = agent_utils.build_mlp(
input_dim=embedding_dim,
hidden_dim=hidden_dim,
output_dim=num_experts,
num_layers=num_layers,
)
self.trunk.apply(agent_utils.weight_init)
self.topk = topk
self._softmax = torch.nn.Softmax(dim=1)
def __init__(
self,
pretrained_embedding_cfg: ConfigType,
num_embeddings: int,
embedding_dim: int,
hidden_dim: int,
num_layers: int,
output_dim: int,
):
"""Encode the task into a vector.
Args:
pretrained_embedding_cfg (ConfigType): config for using pretrained
embeddings.
num_embeddings (int): number of elements in the embedding table. This is
used if pretrained embedding is not used.
embedding_dim (int): dimension for the embedding. This is
used if pretrained embedding is not used.
hidden_dim (int): dimension of the hidden layer of the trunk.
num_layers (int): number of layers in the trunk.
output_dim (int): output dimension of the task encoder.
"""
super().__init__()
if pretrained_embedding_cfg.should_use:
with open(pretrained_embedding_cfg.path_to_load_from) as f:
metadata = json.load(f)
ordered_task_list = pretrained_embedding_cfg.ordered_task_list
pretrained_embedding = torch.Tensor(
[metadata[task] for task in ordered_task_list])
assert num_embeddings == pretrained_embedding.shape[0]
pretrained_embedding_dim = pretrained_embedding.shape[1]
pretrained_embedding = nn.Embedding.from_pretrained(
embeddings=pretrained_embedding,
freeze=True,
)
projection_layer = nn.Sequential(
nn.Linear(in_features=pretrained_embedding_dim,
out_features=2 * embedding_dim),
nn.ReLU(),
nn.Linear(in_features=2 * embedding_dim,
out_features=embedding_dim),
nn.ReLU(),
)
projection_layer.apply(agent_utils.weight_init)
self.embedding = nn.Sequential( # type: ignore [call-overload]
pretrained_embedding,
nn.ReLU(),
projection_layer,
)
else:
self.embedding = nn.Sequential(
nn.Embedding(num_embeddings=num_embeddings,
embedding_dim=embedding_dim),
nn.ReLU(),
)
self.embedding.apply(agent_utils.weight_init)
self.trunk = agent_utils.build_mlp(
input_dim=embedding_dim,
hidden_dim=hidden_dim,
output_dim=output_dim,
num_layers=num_layers,
)
self.trunk.apply(agent_utils.weight_init)