def test_sac_algorithm_init(self):
observation_spec = BoundedTensorSpec((10, ))
discrete_action_spec = BoundedTensorSpec((), dtype='int64')
continuous_action_spec = [
BoundedTensorSpec((3, )),
BoundedTensorSpec((10, ))
]
universal_q_network = partial(
QNetwork, preprocessing_combiner=NestConcat())
critic_network = partial(
CriticNetwork, action_preprocessing_combiner=NestConcat())
# q_network instead of critic_network is needed
self.assertRaises(
AssertionError,
SacAlgorithm,
observation_spec=observation_spec,
action_spec=discrete_action_spec,
q_network_cls=None)
sac = SacAlgorithm(
observation_spec=observation_spec,
action_spec=discrete_action_spec,
q_network_cls=QNetwork)
self.assertEqual(sac._act_type, SacActionType.Discrete)
self.assertEqual(sac.train_state_spec.actor, ())
self.assertEqual(sac.train_state_spec.action.actor_network, ())
# critic_network instead of q_network is needed
self.assertRaises(
AssertionError,
SacAlgorithm,
observation_spec=observation_spec,
action_spec=continuous_action_spec,
critic_network_cls=None)
sac = SacAlgorithm(
observation_spec=observation_spec,
action_spec=continuous_action_spec,
critic_network_cls=critic_network)
self.assertEqual(sac._act_type, SacActionType.Continuous)
self.assertEqual(sac.train_state_spec.action.critic, ())
# action_spec order is incorrect
self.assertRaises(
AssertionError,
SacAlgorithm,
observation_spec=observation_spec,
action_spec=(continuous_action_spec, discrete_action_spec),
q_network_cls=universal_q_network)
sac = SacAlgorithm(
observation_spec=observation_spec,
action_spec=(discrete_action_spec, continuous_action_spec),
q_network_cls=universal_q_network)
self.assertEqual(sac._act_type, SacActionType.Mixed)
self.assertEqual(sac.train_state_spec.actor, ())
def test_value_distribution(self, lstm_hidden_size):
input_spec1 = TensorSpec((3, 20, 20))
input_spec2 = TensorSpec((100, ))
conv_layer_params = ((8, 3, 1), (16, 3, 2, 1))
embedding_dim = 100
image = input_spec1.zeros(outer_dims=(1, ))
vector = input_spec2.zeros(outer_dims=(1, ))
network_ctor, state = self._init(lstm_hidden_size)
value_net = network_ctor(
input_tensor_spec=[input_spec1, input_spec2],
input_preprocessors=[
EmbeddingPreprocessor(
input_spec1,
embedding_dim=embedding_dim,
conv_layer_params=conv_layer_params), None
],
preprocessing_combiner=NestConcat())
value, state = value_net([image, vector], state)
self.assertEqual(value_net._processed_input_tensor_spec.shape[0], 200)
self.assertEqual(value_net.output_spec, TensorSpec(()))
# (batch_size,)
self.assertEqual(value.shape, (1, ))
def __init__(self,
train_state_spec,
action_spec,
feature_spec,
hidden_size=256,
num_replicas=1,
dynamics_network: DynamicsNetwork = None,
name="DynamicsLearningAlgorithm"):
"""Create a DynamicsLearningAlgorithm.
Args:
hidden_size (int|tuple): size of hidden layer(s)
dynamics_network (Network): network for predicting the change of
the next feature based on the previous feature and action.
It should accept input with spec of the format
[feature_spec, encoded_action_spec] and output a tensor of the
shape feature_spec. For discrete action case, encoded_action
is a one-hot representation of the action. For continuous
action, encoded action is the original action.
"""
super().__init__(train_state_spec=train_state_spec, name=name)
flat_action_spec = nest.flatten(action_spec)
assert len(flat_action_spec) == 1, "doesn't support nested action_spec"
flat_feature_spec = nest.flatten(feature_spec)
assert len(
flat_feature_spec) == 1, "doesn't support nested feature_spec"
action_spec = flat_action_spec[0]
if action_spec.is_discrete:
self._num_actions = action_spec.maximum - action_spec.minimum + 1
else:
self._num_actions = action_spec.shape[-1]
self._action_spec = action_spec
self._feature_spec = feature_spec
self._num_replicas = num_replicas
if isinstance(hidden_size, int):
hidden_size = (hidden_size, )
if dynamics_network is None:
encoded_action_spec = TensorSpec((self._num_actions, ),
dtype=torch.float32)
dynamics_network = DynamicsNetwork(
name="dynamics_net",
input_tensor_spec=(feature_spec, encoded_action_spec),
preprocessing_combiner=NestConcat(),
fc_layer_params=hidden_size,
output_tensor_spec=flat_feature_spec[0])
if num_replicas > 1:
self._dynamics_network = dynamics_network.make_parallel(
num_replicas)
else:
self._dynamics_network = dynamics_network
def setUp(self):
self._input_spec = [
TensorSpec((3, 20, 20), torch.float32),
TensorSpec((1, 20, 20), torch.float32)
]
self._image = zero_tensor_from_nested_spec(self._input_spec,
batch_size=1)
self._conv_layer_params = ((8, 3, 1), (16, 3, 2, 1))
self._fc_layer_params = (100, )
self._input_preprocessors = [torch.tanh, None]
self._preprocessing_combiner = NestConcat(dim=1)
def test_mixed_actions(self, net_ctor):
obs_spec = TensorSpec((20, ))
action_spec = dict(x=BoundedTensorSpec((), dtype='int64'),
y=BoundedTensorSpec((3, )))
input_preprocessors = dict(x=EmbeddingPreprocessor(action_spec['x'],
embedding_dim=10),
y=None)
net_ctor = functools.partial(
net_ctor, action_input_processors=input_preprocessors)
# doesn't support mixed actions
self.assertRaises(AssertionError, net_ctor, (obs_spec, action_spec))
# ... unless a combiner is specified
net_ctor((obs_spec, action_spec),
action_preprocessing_combiner=NestConcat())
def test_encoding_network_nested_input(self, lstm):
input_spec = dict(a=TensorSpec((3, 80, 80)),
b=[
TensorSpec((80, )),
BoundedTensorSpec((), dtype="int64"),
dict(x=TensorSpec((100, )),
y=TensorSpec((200, )))
])
imgs = common.zero_tensor_from_nested_spec(input_spec, batch_size=1)
input_preprocessors = dict(
a=EmbeddingPreprocessor(input_spec["a"],
conv_layer_params=((1, 2, 2, 0), ),
embedding_dim=100),
b=[
EmbeddingPreprocessor(input_spec["b"][0], embedding_dim=50),
EmbeddingPreprocessor(input_spec["b"][1], embedding_dim=50),
dict(x=None, y=torch.relu)
])
if lstm:
network_ctor = functools.partial(LSTMEncodingNetwork,
hidden_size=(100, ))
else:
network_ctor = EncodingNetwork
network = network_ctor(input_tensor_spec=input_spec,
input_preprocessors=input_preprocessors,
preprocessing_combiner=NestConcat())
output, _ = network(imgs, state=[(torch.zeros((1, 100)), ) * 2])
if lstm:
self.assertEqual(network.output_spec, TensorSpec((100, )))
self.assertEqual(output.size()[-1], 100)
else:
self.assertEqual(len(list(network.parameters())), 4 + 2 + 1)
self.assertEqual(network.output_spec, TensorSpec((500, )))
self.assertEqual(output.size()[-1], 500)
def __init__(self,
observation_spec,
action_spec,
skill_spec,
config: TrainerConfig,
skill_discriminator_ctor=EncodingNetwork,
skill_encoder_ctor=None,
observation_transformer=math_ops.identity,
optimizer=None,
sparse_reward=False,
debug_summaries=False,
num_steps_per_skill=3,
skill_type="state_difference",
name="Discriminator"):
"""If ``sparse_reward=True``, then the discriminator will only predict
at the skill switching steps.
"""
if skill_spec.is_discrete:
assert isinstance(skill_spec, BoundedTensorSpec)
skill_dim = skill_spec.maximum - skill_spec.minimum + 1
else:
assert len(
skill_spec.shape) == 1, "Only 1D skill vector is supported"
skill_dim = skill_spec.shape[0]
supported_skill_types = [
"state_concatenation",
"state_difference",
"state",
"action",
"action_difference",
"state_action",
"action_concatenation",
]
assert skill_type in supported_skill_types, (
"Skill type must be in: %s" % supported_skill_types)
self._skill_type = skill_type
subtrajectory_spec = get_subtrajectory_spec(num_steps_per_skill,
observation_spec,
action_spec)
if skill_type == "state_concatenation":
discriminator_spec = flatten(subtrajectory_spec.observation)
elif skill_type == "action_concatenation":
discriminator_spec = flatten(subtrajectory_spec.prev_action)
else:
discriminator_spec = get_discriminator_spec(
skill_type, observation_spec, action_spec)
input_preprocessors, preprocessing_combiner = None, None
if is_action_skill(skill_type):
# first project
input_preprocessors = (None, None)
preprocessing_combiner = NestConcat()
discriminator_spec = (observation_spec, discriminator_spec)
skill_encoder = None
if skill_encoder_ctor is not None:
step_spec = BoundedTensorSpec((),
maximum=num_steps_per_skill,
dtype='int64')
skill_encoder = skill_encoder_ctor(
input_preprocessors=(None,
EmbeddingPreprocessor(
input_tensor_spec=step_spec,
embedding_dim=skill_dim)),
preprocessing_combiner=NestConcat(),
input_tensor_spec=(skill_spec, step_spec))
if input_preprocessors is None:
input_preprocessors = (None, )
discriminator_spec = (discriminator_spec, )
input_preprocessors = input_preprocessors + (EmbeddingPreprocessor(
input_tensor_spec=step_spec, embedding_dim=skill_dim), )
discriminator_spec = discriminator_spec + (step_spec, )
skill_dim = skill_encoder.output_spec.shape[0]
skill_disc_inputs = dict(input_preprocessors=input_preprocessors,
preprocessing_combiner=preprocessing_combiner,
input_tensor_spec=discriminator_spec)
if skill_discriminator_ctor.__name__ == "EncodingNetwork":
skill_disc_inputs["last_layer_size"] = skill_dim
else: # ActorDistributionNetwork
skill_disc_inputs["action_spec"] = skill_spec
skill_discriminator = skill_discriminator_ctor(**skill_disc_inputs)
train_state_spec = DiscriminatorState(
first_observation=observation_spec,
untrans_observation=
observation_spec, # prev untransformed observation diff for pred
subtrajectory=subtrajectory_spec)
super().__init__(train_state_spec=train_state_spec,
predict_state_spec=DiscriminatorState(
subtrajectory=subtrajectory_spec,
first_observation=observation_spec),
config=config,
optimizer=optimizer,
debug_summaries=debug_summaries,
name=name)
self._skill_discriminator = skill_discriminator
self._skill_encoder = skill_encoder
# exp observation won't be automatically transformed when it's sampled
# from the replay buffer. We will do this manually.
self._observation_transformer = observation_transformer
self._num_steps_per_skill = num_steps_per_skill
self._sparse_reward = sparse_reward
self._skill_dim = skill_dim
self._high_rl = None
def __init__(self,
x_spec,
y_spec,
model=None,
fc_layers=(256, ),
sampler='buffer',
buffer_size=65536,
optimizer: torch.optim.Optimizer = None,
estimator_type='DV',
averager: EMAverager = None,
name="MIEstimator"):
"""
Args:
x_spec (nested TensorSpec): spec of ``x``
y_spec (nested TensorSpec): spec of ``y``
model (Network): can be called as ``model([x, y])`` and return a Tensor
with ``shape=[batch_size, 1]``. If None, a default MLP with
``fc_layers`` will be created.
fc_layers (tuple[int]): size of hidden layers. Only used if model is
None.
sampler (str): type of sampler used to get samples from marginal
distribution, should be one of ``['buffer', 'double_buffer',
'shuffle', 'shift']``.
buffer_size (int): capacity of buffer for storing y for sampler
'buffer' and 'double_buffer'.
optimzer (torch.optim.Optimzer): optimizer
estimator_type (str): one of 'DV', 'KLD' or 'JSD'
averager (EMAverager): averager used to maintain a moving average
of :math:`exp(T)`. Only used for 'DV' estimator. If None,
a ScalarAdaptiveAverager will be created.
name (str): name of this estimator
"""
assert estimator_type in ['ML', 'DV', 'KLD', 'JSD'
], "Wrong estimator_type %s" % estimator_type
super().__init__(train_state_spec=(), optimizer=optimizer, name=name)
self._x_spec = x_spec
self._y_spec = y_spec
if model is None:
if estimator_type == 'ML':
model = EncodingNetwork(
name="MIEstimator",
input_tensor_spec=x_spec,
fc_layer_params=fc_layers,
preprocessing_combiner=NestConcat(dim=-1))
else:
model = EncodingNetwork(
name="MIEstimator",
input_tensor_spec=[x_spec, y_spec],
preprocessing_combiner=NestConcat(dim=-1),
fc_layer_params=fc_layers,
last_layer_size=1,
last_activation=math_ops.identity)
self._model = model
self._type = estimator_type
if sampler == 'buffer':
self._y_buffer = DataBuffer(y_spec, capacity=buffer_size)
self._sampler = self._buffer_sampler
elif sampler == 'double_buffer':
self._x_buffer = DataBuffer(x_spec, capacity=buffer_size)
self._y_buffer = DataBuffer(y_spec, capacity=buffer_size)
self._sampler = self._double_buffer_sampler
elif sampler == 'shuffle':
self._sampler = self._shuffle_sampler
elif sampler == 'shift':
self._sampler = self._shift_sampler
else:
raise TypeError("Wrong type for sampler %s" % sampler)
if estimator_type == 'DV':
if averager is None:
averager = ScalarAdaptiveAverager()
self._mean_averager = averager
if estimator_type == 'ML':
assert isinstance(
y_spec,
alf.TensorSpec), ("Currently, 'ML' does "
"not support nested y_spec: %s" % y_spec)
assert y_spec.is_continuous, ("Currently, 'ML' does "
"not support discreted y_spec: %s" %
y_spec)
hidden_size = self._model.output_spec.shape[-1]
self._delta_loc_layer = alf.layers.FC(
hidden_size,
y_spec.shape[-1],
kernel_initializer=torch.nn.init.zeros_,
bias_init_value=0.0)
self._delta_scale_layer = alf.layers.FC(
hidden_size,
y_spec.shape[-1],
kernel_initializer=torch.nn.init.zeros_,
bias_init_value=math.log(math.e - 1))
def test_sac_algorithm_mixed(self, use_parallel_network):
num_env = 1
config = TrainerConfig(
root_dir="dummy",
unroll_length=1,
mini_batch_length=2,
mini_batch_size=64,
initial_collect_steps=500,
whole_replay_buffer_training=False,
clear_replay_buffer=False,
num_envs=num_env,
)
env_class = MixedPolicyUnittestEnv
steps_per_episode = 13
env = env_class(num_env, steps_per_episode)
eval_env = env_class(100, steps_per_episode)
obs_spec = env._observation_spec
action_spec = env._action_spec
fc_layer_params = (10, 10, 10)
continuous_projection_net_ctor = partial(
alf.networks.NormalProjectionNetwork,
state_dependent_std=True,
scale_distribution=True,
std_transform=clipped_exp)
actor_network = partial(
ActorDistributionNetwork,
fc_layer_params=fc_layer_params,
continuous_projection_net_ctor=continuous_projection_net_ctor)
q_network = partial(QNetwork,
preprocessing_combiner=NestConcat(),
fc_layer_params=fc_layer_params)
alg2 = SacAlgorithm(observation_spec=obs_spec,
action_spec=action_spec,
actor_network_cls=actor_network,
q_network_cls=q_network,
use_parallel_network=use_parallel_network,
env=env,
config=config,
actor_optimizer=alf.optimizers.Adam(lr=1e-2),
critic_optimizer=alf.optimizers.Adam(lr=1e-2),
alpha_optimizer=alf.optimizers.Adam(lr=1e-2),
debug_summaries=False,
name="MySAC")
eval_env.reset()
for i in range(700):
alg2.train_iter()
if i < config.initial_collect_steps:
continue
eval_env.reset()
eval_time_step = unroll(eval_env, alg2, steps_per_episode - 1)
logging.log_every_n_seconds(
logging.INFO,
"%d reward=%f" % (i, float(eval_time_step.reward.mean())),
n_seconds=1)
self.assertAlmostEqual(1.0,
float(eval_time_step.reward.mean()),
delta=0.2)
def test_nest_selective_concat_specs(self, mask, expected):
ntuple = NTuple(
a=dict(x=TensorSpec((2, 3)), y=TensorSpec((2, 4))),
b=TensorSpec((2, 10)))
ret = NestConcat(mask)(ntuple)
self.assertEqual(ret, expected)
def test_nest_concat_specs(self):
ntuple = NTuple(
a=dict(x=TensorSpec((2, 3)), y=TensorSpec((2, 4))),
b=TensorSpec((2, 10)))
ret = NestConcat()(ntuple)
self.assertEqual(ret, TensorSpec((2, 17)))
def test_nest_concat_tensors(self):
ntuple = NTuple(
a=dict(x=torch.zeros((2, 3)), y=torch.zeros((2, 4))),
b=torch.zeros((2, 10)))
ret = NestConcat()(ntuple)
self.assertTensorEqual(ret, torch.zeros((2, 17)))
def __init__(self,
action_spec,
observation_spec=None,
hidden_size=256,
reward_adapt_speed=8.0,
encoding_net: EncodingNetwork = None,
forward_net: EncodingNetwork = None,
inverse_net: EncodingNetwork = None,
activation=torch.relu_,
optimizer=None,
name="ICMAlgorithm"):
"""Create an ICMAlgorithm.
Args
action_spec (nested TensorSpec): agent's action spec
observation_spec (nested TensorSpec): agent's observation spec. If
not None, then a normalizer will be used to normalize the
observation.
hidden_size (int or tuple[int]): size of hidden layer(s)
reward_adapt_speed (float): how fast to adapt the reward normalizer.
rouphly speaking, the statistics for the normalization is
calculated mostly based on the most recent T/speed samples,
where T is the total number of samples.
encoding_net (Network): network for encoding observation into a
latent feature. Its input is same as the input of this algorithm.
forward_net (Network): network for predicting next feature based on
previous feature and action. It should accept input with spec
[feature_spec, encoded_action_spec] and output a tensor of shape
feature_spec. For discrete action, encoded_action is an one-hot
representation of the action. For continuous action, encoded
action is same as the original action.
inverse_net (Network): network for predicting previous action given
the previous feature and current feature. It should accept input
with spec [feature_spec, feature_spec] and output tensor of
shape (num_actions,).
activation (torch.nn.functional): activation used for constructing
any of the forward net and inverse net, if not provided.
optimizer (torch.optim.Optimizer): The optimizer for training
name (str):
"""
if encoding_net is not None:
feature_spec = encoding_net.output_spec
else:
feature_spec = observation_spec
super(ICMAlgorithm, self).__init__(
train_state_spec=feature_spec,
predict_state_spec=(),
optimizer=optimizer,
name=name)
flat_action_spec = alf.nest.flatten(action_spec)
assert len(
flat_action_spec) == 1, "ICM doesn't suport nested action_spec"
flat_feature_spec = alf.nest.flatten(feature_spec)
assert len(
flat_feature_spec) == 1, "ICM doesn't support nested feature_spec"
action_spec = flat_action_spec[0]
if action_spec.is_discrete:
self._num_actions = int(action_spec.maximum - action_spec.minimum +
1)
else:
self._num_actions = action_spec.shape[-1]
self._action_spec = action_spec
self._observation_normalizer = None
if observation_spec is not None:
self._observation_normalizer = AdaptiveNormalizer(
tensor_spec=observation_spec)
feature_dim = flat_feature_spec[0].shape[-1]
self._encoding_net = encoding_net
if isinstance(hidden_size, int):
hidden_size = (hidden_size, )
if forward_net is None:
encoded_action_spec = TensorSpec((self._num_actions, ),
dtype=torch.float32)
forward_net = EncodingNetwork(
name="forward_net",
input_tensor_spec=[feature_spec, encoded_action_spec],
preprocessing_combiner=NestConcat(),
fc_layer_params=hidden_size,
activation=activation,
last_layer_size=feature_dim,
last_activation=math_ops.identity)
self._forward_net = forward_net
if inverse_net is None:
inverse_net = EncodingNetwork(
name="inverse_net",
input_tensor_spec=[feature_spec, feature_spec],
preprocessing_combiner=NestConcat(),
fc_layer_params=hidden_size,
activation=activation,
last_layer_size=self._num_actions,
last_activation=math_ops.identity,
last_kernel_initializer=torch.nn.init.zeros_)
self._inverse_net = inverse_net
self._reward_normalizer = ScalarAdaptiveNormalizer(
speed=reward_adapt_speed)
def test_conditional_vae(self):
"""Test for one dimensional Gaussion, conditioned on a Bernoulli variable.
"""
prior_input_spec = BoundedTensorSpec((), 'int64')
z_prior_network = EncodingNetwork(
TensorSpec(
(prior_input_spec.maximum - prior_input_spec.minimum + 1, )),
fc_layer_params=(10, ) * 2,
last_layer_size=2 * self._latent_dim,
last_activation=math_ops.identity)
preprocess_network = EncodingNetwork(
input_tensor_spec=(
z_prior_network.input_tensor_spec,
self._input_spec,
z_prior_network.output_spec,
),
preprocessing_combiner=NestConcat(),
fc_layer_params=(10, ) * 2,
last_layer_size=self._latent_dim,
last_activation=math_ops.identity)
encoder = vae.VariationalAutoEncoder(
self._latent_dim,
preprocess_network=preprocess_network,
z_prior_network=z_prior_network)
decoding_layers = FC(self._latent_dim, 1)
optimizer = torch.optim.Adam(
list(encoder.parameters()) + list(decoding_layers.parameters()),
lr=0.1)
x_train = self._input_spec.randn(outer_dims=(10000, ))
y_train = x_train.clone()
y_train[:5000] = y_train[:5000] + 1.0
pr_train = torch.cat([
prior_input_spec.zeros(outer_dims=(5000, )),
prior_input_spec.ones(outer_dims=(5000, ))
],
dim=0)
x_test = self._input_spec.randn(outer_dims=(100, ))
y_test = x_test.clone()
y_test[:50] = y_test[:50] + 1.0
pr_test = torch.cat([
prior_input_spec.zeros(outer_dims=(50, )),
prior_input_spec.ones(outer_dims=(50, ))
],
dim=0)
pr_test = torch.nn.functional.one_hot(
pr_test,
int(z_prior_network.input_tensor_spec.shape[0])).to(torch.float32)
for _ in range(self._epochs):
idx = torch.randperm(x_train.shape[0])
x_train = x_train[idx]
y_train = y_train[idx]
pr_train = pr_train[idx]
for i in range(0, x_train.shape[0], self._batch_size):
optimizer.zero_grad()
batch = x_train[i:i + self._batch_size]
y_batch = y_train[i:i + self._batch_size]
pr_batch = torch.nn.functional.one_hot(
pr_train[i:i + self._batch_size],
int(z_prior_network.input_tensor_spec.shape[0])).to(
torch.float32)
alg_step = encoder.train_step([pr_batch, batch])
outputs = decoding_layers(alg_step.output)
loss = torch.mean(100 * self._loss_f(y_batch - outputs) +
alg_step.info.loss)
loss.backward()
optimizer.step()
y_hat_test = decoding_layers(
encoder.train_step([pr_test, x_test]).output)
reconstruction_loss = float(
torch.mean(self._loss_f(y_test - y_hat_test)))
print("reconstruction_loss:", reconstruction_loss)
self.assertLess(reconstruction_loss, 0.05)
def __init__(self,
input_tensor_spec,
memory_size,
core_size,
num_prememory_layers,
num_memory_layers,
num_attention_heads,
d_ff,
centralized_memory=True,
input_preprocessors=None,
name="TransformerNetwork"):
"""
Args:
input_tensor_spec (nested TensorSpec): the (nested) tensor spec of
the input. If ``input_tensor_spec`` is not nested, it should
represent a rank-2 tensor of shape ``[input_size, d_model]``, where
``input_size`` is the length of the input sequence, and ``d_model``
is the dimension of embedding.
memory_size (int): size of memory.
core_size (int): size of core (i.e. number of embeddings of core)
num_prememory_layers (int): number of TransformerBlock calculation
without using memory
num_memory_layers (int): number of TransformerBlock calculation
using memory
num_attention_heads (int): number of attention heads for each
``TransformerBlock``
d_ff (int): the size of the hidden layer of the feedforward network
in each ``TransformerBlock``
centralized_memory (bool): if False, there will be a separate memory
for each memory layers. if True, there will be a single memory
for all the memroy layers and it is updated using the last core
embeddings.
input_preprocessors (nested Network|nn.Module): a nest of
preprocessor networks, each of which will be applied to the
corresponding input. If not None, then it must have the same
structure with ``input_tensor_spec``. If any element is None, then
it will be treated as math_ops.identity. This arg is helpful if
you want to have separate preprocessings for different inputs by
configuring a gin file without changing the code. For example,
embedding a discrete input before concatenating it to another
continuous vector. The output_spec of each input preprocessor i
should be [input_size_i, d_model]. The result of all the preprocessors
will be concatenated as a Tensor of shape ``[batch_size, input_size, d_model]``,
where ``input_size = sum_i input_size_i``.
"""
preprocessing_combiner = None
if input_preprocessors is not None:
preprocessing_combiner = NestConcat(dim=-2)
super().__init__(input_tensor_spec,
input_preprocessors,
preprocessing_combiner=preprocessing_combiner,
name=name)
assert self._processed_input_tensor_spec.ndim == 2
input_size, d_model = self._processed_input_tensor_spec.shape
if centralized_memory:
self._memories = [FIFOMemory(d_model, memory_size)]
else:
self._memories = [
FIFOMemory(d_model, memory_size)
for _ in range(num_memory_layers)
]
self._centralized_memory = centralized_memory
self._core_size = core_size
self._core_embedding = nn.Parameter(torch.Tensor(
1, core_size, d_model))
nn.init.uniform_(self._core_embedding, -0.1, 0.1)
self._state_spec = [mem.state_spec for mem in self._memories]
self._num_memory_layers = num_memory_layers
self._num_prememory_layers = num_prememory_layers
self._transformers = nn.ModuleList()
for i in range(num_prememory_layers):
self._transformers.append(
alf.layers.TransformerBlock(d_model=d_model,
num_heads=num_attention_heads,
memory_size=input_size + core_size,
positional_encoding='abs'))
for i in range(num_memory_layers):
self._transformers.append(
alf.layers.TransformerBlock(d_model=d_model,
num_heads=num_attention_heads,
memory_size=memory_size +
input_size + core_size,
positional_encoding='abs'))
def test_input_preprocessor(self, lstm, preproc):
def _check_with_shared_param(net1, net2, shared_subnet=None):
net1_params = set(net1.parameters())
net2_params = set(net2.parameters())
# check that net1 and net2 share paramsters with shared_subnet
if shared_subnet is not None:
shared_params = set(shared_subnet.parameters())
for p in shared_params:
self.assertTrue((p in net1_params) and (p in net2_params))
# for the rest part, net1 and net2 do not share parameters
for p1, p2 in zip(net1_params, net2_params):
if shared_subnet is None or p1 not in shared_params:
self.assertTrue(p1 is not p2)
# 1) test input_preprocessor copy and each copy has its own parameters
input_preprocessor = preproc
input_preprocessor_copy = input_preprocessor.copy()
if not preproc._singleton_instance:
_check_with_shared_param(input_preprocessor,
input_preprocessor_copy)
elif preproc._singleton_instance:
_check_with_shared_param(input_preprocessor,
input_preprocessor_copy,
input_preprocessor)
if lstm:
network_ctor = functools.partial(LSTMEncodingNetwork,
hidden_size=(1, ),
post_fc_layer_params=(2, 2))
else:
network_ctor = functools.partial(EncodingNetwork,
fc_layer_params=(10, 10))
net = network_ctor(
input_tensor_spec=[
TestInputpreprocessor.input_spec,
TestInputpreprocessor.input_spec
],
input_preprocessors=[input_preprocessor, torch.relu],
preprocessing_combiner=NestConcat(dim=1))
# 2) test copied network has its own parameters, including
# parameters from input preprocessors
copied_net = net.copy()
if not preproc._singleton_instance:
_check_with_shared_param(net, copied_net)
else:
_check_with_shared_param(net, copied_net, input_preprocessor)
# 3) test for each replica of the NaiveParallelNetwork has its own
# parameters, including parameters from input preprocessors
replicas = 2
p_net = alf.networks.network.NaiveParallelNetwork(net, replicas)
if not preproc._singleton_instance:
_check_with_shared_param(p_net._networks[0], p_net._networks[1])
else:
_check_with_shared_param(p_net._networks[0], p_net._networks[1],
input_preprocessor)
# 4) test network forward
batch_size = 6
batch = TestInputpreprocessor.input_spec.zeros(
outer_dims=(batch_size, ))
if lstm:
state = [(torch.zeros((batch_size, 1)), ) * 2]
p_state = [(torch.zeros((batch_size, replicas, 1)), ) * 2]
else:
state = ()
p_state = ()
net([batch, batch], state)
p_net([batch, batch], p_state)
def test_nest_selective_concat_tensors(self, mask, expected):
ntuple = NTuple(
a=dict(x=torch.zeros((2, 3)), y=torch.zeros((2, 4))),
b=torch.zeros((2, 10)))
ret = NestConcat(mask)(ntuple)
self.assertTensorEqual(ret, expected)
