def __init__(self, specs: BehaviorSpec,
settings: CuriositySettings) -> None:
super().__init__()
self._policy_specs = specs
state_encoder_settings = NetworkSettings(
normalize=False,
hidden_units=settings.encoding_size,
num_layers=2,
vis_encode_type=EncoderType.SIMPLE,
memory=None,
)
self._state_encoder = NetworkBody(specs.observation_shapes,
state_encoder_settings)
self._action_flattener = ModelUtils.ActionFlattener(specs)
self.inverse_model_action_prediction = torch.nn.Sequential(
LinearEncoder(2 * settings.encoding_size, 1, 256),
linear_layer(256, self._action_flattener.flattened_size),
)
self.forward_model_next_state_prediction = torch.nn.Sequential(
LinearEncoder(
settings.encoding_size + self._action_flattener.flattened_size,
1, 256),
linear_layer(256, settings.encoding_size),
)
def __init__(
self,
x_self_size: int,
entities_sizes: List[int],
embedding_size: int,
output_size: Optional[int] = None,
):
super().__init__()
self.self_size = x_self_size
self.entities_sizes = entities_sizes
self.entities_num_max_elements: Optional[List[int]] = None
self.ent_encoders = torch.nn.ModuleList([
LinearEncoder(self.self_size + ent_size, 2, embedding_size)
for ent_size in self.entities_sizes
])
self.attention = MultiHeadAttention(
query_size=embedding_size,
key_size=embedding_size,
value_size=embedding_size,
output_size=embedding_size,
num_heads=4,
embedding_size=embedding_size,
)
self.residual_layer = LinearEncoder(embedding_size, 1, embedding_size)
if output_size is None:
output_size = embedding_size
self.x_self_residual_layer = LinearEncoder(
embedding_size + x_self_size, 1, output_size)
def __init__(self, specs: BehaviorSpec,
settings: CuriositySettings) -> None:
super().__init__()
self._action_spec = specs.action_spec
state_encoder_settings = NetworkSettings(
normalize=False,
hidden_units=settings.encoding_size,
num_layers=2,
vis_encode_type=EncoderType.SIMPLE,
memory=None,
)
self._state_encoder = NetworkBody(specs.observation_specs,
state_encoder_settings)
self._action_flattener = ActionFlattener(self._action_spec)
self.inverse_model_action_encoding = torch.nn.Sequential(
LinearEncoder(2 * settings.encoding_size, 1, 256))
if self._action_spec.continuous_size > 0:
self.continuous_action_prediction = linear_layer(
256, self._action_spec.continuous_size)
if self._action_spec.discrete_size > 0:
self.discrete_action_prediction = linear_layer(
256, sum(self._action_spec.discrete_branches))
self.forward_model_next_state_prediction = torch.nn.Sequential(
LinearEncoder(
settings.encoding_size + self._action_flattener.flattened_size,
1, 256),
linear_layer(256, settings.encoding_size),
)
def __init__(
self,
observation_specs: List[ObservationSpec],
network_settings: NetworkSettings,
encoded_act_size: int = 0,
):
super().__init__()
self.normalize = network_settings.normalize
self.use_lstm = network_settings.memory is not None
self.h_size = network_settings.hidden_units
self.m_size = (network_settings.memory.memory_size
if network_settings.memory is not None else 0)
self.processors, self.embedding_sizes = ModelUtils.create_input_processors(
observation_specs,
self.h_size,
network_settings.vis_encode_type,
normalize=self.normalize,
)
entity_num_max: int = 0
var_processors = [
p for p in self.processors if isinstance(p, EntityEmbedding)
]
for processor in var_processors:
entity_max: int = processor.entity_num_max_elements
# Only adds entity max if it was known at construction
if entity_max > 0:
entity_num_max += entity_max
if len(var_processors) > 0:
if sum(self.embedding_sizes):
self.x_self_encoder = LinearEncoder(
sum(self.embedding_sizes),
1,
self.h_size,
kernel_init=Initialization.Normal,
kernel_gain=(0.125 / self.h_size)**0.5,
)
self.rsa = ResidualSelfAttention(self.h_size, entity_num_max)
total_enc_size = sum(self.embedding_sizes) + self.h_size
else:
total_enc_size = sum(self.embedding_sizes)
total_enc_size += encoded_act_size
self.linear_encoder = LinearEncoder(total_enc_size,
network_settings.num_layers,
self.h_size)
if self.use_lstm:
self.lstm = LSTM(self.h_size, self.m_size)
else:
self.lstm = None # type: ignore
def __init__(
self,
observation_shapes: List[Tuple[int, ...]],
network_settings: NetworkSettings,
encoded_act_size: int = 0,
):
super().__init__()
self.normalize = network_settings.normalize
self.use_lstm = network_settings.memory is not None
self.h_size = network_settings.hidden_units
self.m_size = (network_settings.memory.memory_size
if network_settings.memory is not None else 0)
self.visual_processors, self.vector_processors, encoder_input_size = ModelUtils.create_input_processors(
observation_shapes,
self.h_size,
network_settings.vis_encode_type,
normalize=self.normalize,
)
total_enc_size = encoder_input_size + encoded_act_size
self.linear_encoder = LinearEncoder(total_enc_size,
network_settings.num_layers,
self.h_size)
if self.use_lstm:
self.lstm = LSTM(self.h_size, self.m_size)
else:
self.lstm = None # type: ignore
def __init__(
self,
entity_size: int,
entity_num_max_elements: Optional[int],
embedding_size: int,
):
"""
Constructs an EntityEmbedding module.
:param x_self_size: Size of "self" entity.
:param entity_size: Size of other entities.
:param entity_num_max_elements: Maximum elements for a given entity, None for unrestricted.
Needs to be assigned in order for model to be exportable to ONNX and Barracuda.
:param embedding_size: Embedding size for the entity encoder.
:param concat_self: Whether to concatenate x_self to entities. Set True for ego-centric
self-attention.
"""
super().__init__()
self.self_size: int = 0
self.entity_size: int = entity_size
self.entity_num_max_elements: int = -1
if entity_num_max_elements is not None:
self.entity_num_max_elements = entity_num_max_elements
self.embedding_size = embedding_size
# Initialization scheme from http://www.cs.toronto.edu/~mvolkovs/ICML2020_tfixup.pdf
self.self_ent_encoder = LinearEncoder(
self.entity_size,
1,
self.embedding_size,
kernel_init=Initialization.Normal,
kernel_gain=(0.125 / self.embedding_size)**0.5,
)
def __init__(
self,
sensor_specs: List[SensorSpec],
network_settings: NetworkSettings,
encoded_act_size: int = 0,
):
super().__init__()
self.normalize = network_settings.normalize
self.use_lstm = network_settings.memory is not None
self.h_size = network_settings.hidden_units
self.m_size = (network_settings.memory.memory_size
if network_settings.memory is not None else 0)
self.processors, self.embedding_sizes = ModelUtils.create_input_processors(
sensor_specs,
self.h_size,
network_settings.vis_encode_type,
normalize=self.normalize,
)
total_enc_size = sum(self.embedding_sizes) + encoded_act_size
self.linear_encoder = LinearEncoder(total_enc_size,
network_settings.num_layers,
self.h_size)
if self.use_lstm:
self.lstm = LSTM(self.h_size, self.m_size)
else:
self.lstm = None # type: ignore
def __init__(
self,
observation_specs: List[ObservationSpec],
network_settings: NetworkSettings,
encoded_act_size: int = 0,
):
super().__init__()
self.normalize = network_settings.normalize
self.use_lstm = network_settings.memory is not None
self.h_size = network_settings.hidden_units
self.m_size = (network_settings.memory.memory_size
if network_settings.memory is not None else 0)
self.observation_encoder = ObservationEncoder(
observation_specs,
self.h_size,
network_settings.vis_encode_type,
self.normalize,
)
self.processors = self.observation_encoder.processors
total_enc_size = self.observation_encoder.total_enc_size
total_enc_size += encoded_act_size
self.linear_encoder = LinearEncoder(total_enc_size,
network_settings.num_layers,
self.h_size)
if self.use_lstm:
self.lstm = LSTM(self.h_size, self.m_size)
else:
self.lstm = None # type: ignore
def create_residual_self_attention(
input_processors: nn.ModuleList, embedding_sizes: List[int],
hidden_size: int
) -> Tuple[Optional[ResidualSelfAttention], Optional[LinearEncoder]]:
"""
Creates an RSA if there are variable length observations found in the input processors.
:param input_processors: A ModuleList of input processors as returned by the function
create_input_processors().
:param embedding sizes: A List of embedding sizes as returned by create_input_processors().
:param hidden_size: The hidden size to use for the RSA.
:returns: A Tuple of the RSA itself, a self encoder, and the embedding size after the RSA.
Returns None for the RSA and encoder if no var len inputs are detected.
"""
rsa, x_self_encoder = None, None
entity_num_max: int = 0
var_processors = [
p for p in input_processors if isinstance(p, EntityEmbedding)
]
for processor in var_processors:
entity_max: int = processor.entity_num_max_elements
# Only adds entity max if it was known at construction
if entity_max > 0:
entity_num_max += entity_max
if len(var_processors) > 0:
if sum(embedding_sizes):
x_self_encoder = LinearEncoder(
sum(embedding_sizes),
1,
hidden_size,
kernel_init=Initialization.Normal,
kernel_gain=(0.125 / hidden_size)**0.5,
)
rsa = ResidualSelfAttention(hidden_size, entity_num_max)
return rsa, x_self_encoder
def test_predict_minimum_training():
# of 5 numbers, predict index of min
np.random.seed(1336)
torch.manual_seed(1336)
n_k = 5
size = n_k + 1
embedding_size = 64
entity_embeddings = EntityEmbeddings(size, [size],
embedding_size, [n_k],
concat_self=False)
transformer = ResidualSelfAttention(embedding_size)
l_layer = LinearEncoder(embedding_size, 2, n_k)
loss = torch.nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(
list(entity_embeddings.parameters()) + list(transformer.parameters()) +
list(l_layer.parameters()),
lr=0.001,
weight_decay=1e-6,
)
batch_size = 200
onehots = ModelUtils.actions_to_onehot(
torch.range(0, n_k - 1).unsqueeze(1), [n_k])[0]
onehots = onehots.expand((batch_size, -1, -1))
losses = []
for _ in range(400):
num = np.random.randint(0, n_k)
inp = torch.rand((batch_size, num + 1, 1))
with torch.no_grad():
# create the target : The minimum
argmin = torch.argmin(inp, dim=1)
argmin = argmin.squeeze()
argmin = argmin.detach()
sliced_oh = onehots[:, :num + 1]
inp = torch.cat([inp, sliced_oh], dim=2)
embeddings = entity_embeddings(inp, [inp])
masks = EntityEmbeddings.get_masks([inp])
prediction = transformer(embeddings, masks)
prediction = l_layer(prediction)
ce = loss(prediction, argmin)
losses.append(ce.item())
print(ce.item())
optimizer.zero_grad()
ce.backward()
optimizer.step()
assert np.array(losses[-20:]).mean() < 0.1
def add_self_embedding(self, size: int) -> None:
self.self_size = size
self.self_ent_encoder = LinearEncoder(
self.self_size + self.entity_size,
1,
self.embedding_size,
kernel_init=Initialization.Normal,
kernel_gain=(0.125 / self.embedding_size)**0.5,
)
def __init__(self, specs: BehaviorSpec,
settings: CuriositySettings) -> None:
super().__init__()
self._action_spec = specs.action_spec
state_encoder_settings = settings.network_settings
if state_encoder_settings.memory is not None:
state_encoder_settings.memory = None
logger.warning(
"memory was specified in network_settings but is not supported by Curiosity. It is being ignored."
)
self._state_encoder = NetworkBody(specs.observation_specs,
state_encoder_settings)
self._action_flattener = ActionFlattener(self._action_spec)
self.inverse_model_action_encoding = torch.nn.Sequential(
LinearEncoder(2 * state_encoder_settings.hidden_units, 1, 256))
if self._action_spec.continuous_size > 0:
self.continuous_action_prediction = linear_layer(
256, self._action_spec.continuous_size)
if self._action_spec.discrete_size > 0:
self.discrete_action_prediction = linear_layer(
256, sum(self._action_spec.discrete_branches))
self.forward_model_next_state_prediction = torch.nn.Sequential(
LinearEncoder(
state_encoder_settings.hidden_units +
self._action_flattener.flattened_size,
1,
256,
),
linear_layer(256, state_encoder_settings.hidden_units),
)
def __init__(
self,
observation_specs: List[ObservationSpec],
network_settings: NetworkSettings,
action_spec: ActionSpec,
):
super().__init__()
self.normalize = network_settings.normalize
self.use_lstm = network_settings.memory is not None
self.h_size = network_settings.hidden_units
self.m_size = (network_settings.memory.memory_size
if network_settings.memory is not None else 0)
self.action_spec = action_spec
self.observation_encoder = ObservationEncoder(
observation_specs,
self.h_size,
network_settings.vis_encode_type,
self.normalize,
)
self.processors = self.observation_encoder.processors
# Modules for multi-agent self-attention
obs_only_ent_size = self.observation_encoder.total_enc_size
q_ent_size = (obs_only_ent_size +
sum(self.action_spec.discrete_branches) +
self.action_spec.continuous_size)
attention_embeding_size = self.h_size
self.obs_encoder = EntityEmbedding(obs_only_ent_size, None,
attention_embeding_size)
self.obs_action_encoder = EntityEmbedding(q_ent_size, None,
attention_embeding_size)
self.self_attn = ResidualSelfAttention(attention_embeding_size)
self.linear_encoder = LinearEncoder(
attention_embeding_size,
network_settings.num_layers,
self.h_size,
kernel_gain=(0.125 / self.h_size)**0.5,
)
if self.use_lstm:
self.lstm = LSTM(self.h_size, self.m_size)
else:
self.lstm = None # type: ignore
self._current_max_agents = torch.nn.Parameter(torch.as_tensor(1),
requires_grad=False)
def __init__(
self,
x_self_size: int,
entity_sizes: List[int],
entity_num_max_elements: List[int],
embedding_size: int,
concat_self: bool = True,
):
super().__init__()
self.self_size: int = x_self_size
self.entity_sizes: List[int] = entity_sizes
self.entity_num_max_elements: List[int] = entity_num_max_elements
self.concat_self: bool = concat_self
# If not concatenating self, input to encoder is just entity size
if not concat_self:
self.self_size = 0
self.ent_encoders = torch.nn.ModuleList([
LinearEncoder(self.self_size + ent_size, 1, embedding_size)
for ent_size in self.entity_sizes
])
def __init__(
self,
x_self_size: int,
entity_sizes: List[int],
embedding_size: int,
entity_num_max_elements: Optional[List[int]] = None,
concat_self: bool = True,
):
"""
Constructs an EntityEmbeddings module.
:param x_self_size: Size of "self" entity.
:param entity_sizes: List of sizes for other entities. Should be of length
equivalent to the number of entities.
:param embedding_size: Embedding size for entity encoders.
:param entity_num_max_elements: Maximum elements in an entity, None for unrestricted.
Needs to be assigned in order for model to be exportable to ONNX and Barracuda.
:param concat_self: Whether to concatenate x_self to entites. Set True for ego-centric
self-attention.
"""
super().__init__()
self.self_size: int = x_self_size
self.entity_sizes: List[int] = entity_sizes
self.entity_num_max_elements: List[int] = [-1] * len(entity_sizes)
if entity_num_max_elements is not None:
self.entity_num_max_elements = entity_num_max_elements
self.concat_self: bool = concat_self
# If not concatenating self, input to encoder is just entity size
if not concat_self:
self.self_size = 0
# Initialization scheme from http://www.cs.toronto.edu/~mvolkovs/ICML2020_tfixup.pdf
self.ent_encoders = torch.nn.ModuleList([
LinearEncoder(
self.self_size + ent_size,
1,
embedding_size,
kernel_init=Initialization.Normal,
kernel_gain=(0.125 / embedding_size)**0.5,
) for ent_size in self.entity_sizes
])
self.embedding_norm = LayerNorm()
