def test_concat_block():
""" perception test """
for d0, d1 in [([100, 16], [100, 16]), ([100, 1, 16], [100, 1, 16])]:
in_dict = build_multi_input_dict(dims=[d0, d1])
net = ConcatenationBlock(in_keys=["in_key_0", "in_key_1"],
out_keys="concat",
in_shapes=[d0[1:], d1[1:]],
concat_dim=-1)
str(net)
out_dict = net(in_dict)
assert isinstance(out_dict, Dict)
assert out_dict["concat"].shape[-1] == 32
assert net.out_shapes() == [out_dict["concat"].shape[1:]]
def __init__(self, obs_shapes: Dict[str, Sequence[int]],
non_lin: Union[str, type(nn.Module)], hidden_units: List[int]):
self.obs_shapes = obs_shapes
# Maze relies on dictionaries to represent the inference graph
self.perception_dict = OrderedDict()
# build latent feature embedding block
self.perception_dict['latent_inventory'] = DenseBlock(
in_keys='observation_inventory', out_keys='latent_inventory', in_shapes=obs_shapes['observation_inventory'],
hidden_units=[128], non_lin=non_lin)
# build latent pixel embedding block
self.perception_dict['latent_screen'] = VGGConvolutionDenseBlock(
in_keys='observation_screen', out_keys='latent_screen', in_shapes=obs_shapes['observation_screen'],
non_lin=non_lin, hidden_channels=[8, 16, 32], hidden_units=[32])
# Concatenate latent features
self.perception_dict['latent_concat'] = ConcatenationBlock(
in_keys=['latent_inventory', 'latent_screen'], out_keys='latent_concat',
in_shapes=self.perception_dict['latent_inventory'].out_shapes() +
self.perception_dict['latent_screen'].out_shapes(), concat_dim=-1)
# Add latent dense block
self.perception_dict['latent_dense'] = DenseBlock(
in_keys='latent_concat', out_keys='latent_dense', hidden_units=hidden_units, non_lin=non_lin,
in_shapes=self.perception_dict['latent_concat'].out_shapes()
)
# Add recurrent block
self.perception_dict['latent'] = LSTMLastStepBlock(
in_keys='latent_dense', out_keys='latent', in_shapes=self.perception_dict['latent_dense'].out_shapes(),
hidden_size=32, num_layers=1, bidirectional=False, non_lin=non_lin
)
def __init__(self, obs_shapes: Dict[str, Sequence[int]], non_lin: Union[str, type(nn.Module)]):
nn.Module.__init__(self)
self.obs_shapes = obs_shapes
hidden_units = 32
self.perception_dict = OrderedDict()
self.perception_dict['order_feat'] = DenseBlock(
in_keys='ordered_piece', out_keys='order_feat', in_shapes=self.obs_shapes['ordered_piece'],
hidden_units=[hidden_units], non_lin=non_lin)
self.perception_dict['selected_feat'] = DenseBlock(
in_keys='selected_piece', out_keys='selected_feat', in_shapes=self.obs_shapes['selected_piece'],
hidden_units=[hidden_units], non_lin=non_lin)
self.perception_dict['latent'] = ConcatenationBlock(
in_keys=['order_feat', 'selected_feat'], out_keys='latent',
in_shapes=[[hidden_units], [hidden_units], [hidden_units]], concat_dim=-1)
self.perception_dict['value'] = LinearOutputBlock(
in_keys='latent', out_keys='value', in_shapes=self.perception_dict['latent'].out_shapes(), output_units=1)
in_keys = ['ordered_piece', 'selected_piece']
self.perception_net = InferenceBlock(
in_keys=in_keys, out_keys='value',
in_shapes=[self.obs_shapes[key] for key in in_keys],
perception_blocks=self.perception_dict)
# initialize model weights
self.perception_net.apply(make_module_init_normc(1.0))
self.perception_dict['value'].apply(make_module_init_normc(0.01))
def __init__(self, obs_shapes: Dict[str, Sequence[int]],
non_lin: Union[str,
type(nn.Module)], hidden_units: List[int]):
nn.Module.__init__(self)
# Maze relies on dictionaries to represent the inference graph
self.perception_dict = OrderedDict()
# build latent feature embedding block
self.perception_dict['latent_inventory'] = DenseBlock(
in_keys='observation_inventory',
out_keys='latent_inventory',
in_shapes=obs_shapes['observation_inventory'],
hidden_units=[128],
non_lin=non_lin)
# Concatenate latent features
self.perception_dict['latent_concat'] = ConcatenationBlock(
in_keys=['latent_inventory', 'latent_screen'],
out_keys='latent_concat',
in_shapes=self.perception_dict['latent_inventory'].out_shapes() +
[obs_shapes['latent_screen']],
concat_dim=-1)
# Add latent dense block
self.perception_dict['latent_dense'] = DenseBlock(
in_keys='latent_concat',
out_keys='latent_dense',
hidden_units=hidden_units,
non_lin=non_lin,
in_shapes=self.perception_dict['latent_concat'].out_shapes())
# Add recurrent block
self.perception_dict['latent'] = LSTMLastStepBlock(
in_keys='latent_dense',
out_keys='latent',
in_shapes=self.perception_dict['latent_dense'].out_shapes(),
hidden_size=32,
num_layers=1,
bidirectional=False,
non_lin=non_lin)
# build action heads
self.perception_dict['value'] = LinearOutputBlock(
in_keys='latent',
out_keys='value',
in_shapes=self.perception_dict['latent'].out_shapes(),
output_units=1)
# build inference block
in_keys = list(obs_shapes.keys())
self.perception_net = InferenceBlock(
in_keys=in_keys,
out_keys='value',
in_shapes=[obs_shapes[key] for key in in_keys],
perception_blocks=self.perception_dict)
# apply weight init
self.perception_net.apply(make_module_init_normc(1.0))
self.perception_dict['value'].apply(make_module_init_normc(0.01))
def __init__(self, obs_shapes: Dict[str, Sequence[int]],
non_lin: Union[str, type(nn.Module)]):
nn.Module.__init__(self)
# initialize the perception dictionary
self.perception_dict = OrderedDict()
# concatenate all observations in dictionary
self.perception_dict['concat'] = ConcatenationBlock(
in_keys=[
'cart_position', 'cart_velocity', 'pole_angle',
'pole_angular_velocity'
],
out_keys='concat',
in_shapes=[
obs_shapes['cart_position'], obs_shapes['cart_velocity'],
obs_shapes['pole_angle'], obs_shapes['pole_angular_velocity']
],
concat_dim=-1)
# process concatenated representation with two dense layers
self.perception_dict['embedding'] = DenseBlock(
in_keys='concat',
in_shapes=self.perception_dict['concat'].out_shapes(),
hidden_units=[128, 128],
non_lin=non_lin,
out_keys='embedding')
# add a linear output block
self.perception_dict['value'] = LinearOutputBlock(
in_keys='embedding',
out_keys='value',
in_shapes=self.perception_dict['embedding'].out_shapes(),
output_units=1)
# compile an inference block
self.perception_net = InferenceBlock(
in_keys=[
'cart_position', 'cart_velocity', 'pole_angle',
'pole_angular_velocity'
],
out_keys='value',
in_shapes=[
obs_shapes[key] for key in [
'cart_position', 'cart_velocity', 'pole_angle',
'pole_angular_velocity'
]
],
perception_blocks=self.perception_dict)
# initialize model weights
self.perception_net.apply(make_module_init_normc(1.0))
self.perception_dict['value'].apply(make_module_init_normc(0.01))
def __init__(self, obs_shapes: Dict[str, Sequence[int]],
action_spaces_dict: Dict[Union[str, int], spaces.Space],
non_lin: Union[str, type(nn.Module)]):
super().__init__()
self.obs_shapes = obs_shapes
# build perception part
self.perception_dict = OrderedDict()
self.perception_dict['latent-obs'] = DenseBlock(
in_keys="observation",
out_keys="latent-obs",
in_shapes=obs_shapes['observation'],
hidden_units=[256],
non_lin=non_lin)
self.perception_dict['latent-act'] = DenseBlock(
in_keys="action",
out_keys="latent-act",
in_shapes=obs_shapes['action'],
hidden_units=[256],
non_lin=non_lin)
self.perception_dict['concat'] = ConcatenationBlock(
in_keys=['latent-obs', 'latent-act'],
in_shapes=self.perception_dict['latent-obs'].out_shapes() +
self.perception_dict['latent-act'].out_shapes(),
concat_dim=-1,
out_keys='concat')
self.perception_dict['latent'] = DenseBlock(
in_keys="concat",
out_keys="latent",
in_shapes=self.perception_dict['concat'].out_shapes(),
hidden_units=[256],
non_lin=non_lin)
# build action head
self.perception_dict['q_value'] = LinearOutputBlock(
in_keys="latent",
out_keys="q_value",
in_shapes=self.perception_dict['latent'].out_shapes(),
output_units=1)
self.perception_net = InferenceBlock(
in_keys=['observation', 'action'],
out_keys='q_value',
in_shapes=[
self.obs_shapes['observation'], self.obs_shapes['action']
],
perception_blocks=self.perception_dict)
# initialize model weights
self.perception_net.apply(make_module_init_normc(1.0))
self.perception_dict['q_value'].apply(make_module_init_normc(0.01))
def __init__(self, obs_shapes: Dict[str, Sequence[int]],
non_lin: type(nn.Module)):
nn.Module.__init__(self)
self.obs_shapes = obs_shapes
perception_dict: Dict[str, PerceptionBlock] = dict()
for in_key, in_shape in self.obs_shapes.items():
if len(in_shape) > 1:
next_in_key = f'{in_key}_flat'
perception_dict[next_in_key] = FlattenBlock(
in_keys=in_key,
in_shapes=in_shape,
out_keys=next_in_key,
num_flatten_dims=len(in_shape))
next_in_shape = perception_dict[next_in_key].out_shapes()
else:
next_in_key = in_key
next_in_shape = in_shape
perception_dict[f'{in_key}_encoded_feat'] = DenseBlock(
in_keys=next_in_key,
in_shapes=next_in_shape,
out_keys=f'{in_key}_encoded_feat',
non_lin=non_lin,
hidden_units=[16])
perception_dict[f'{in_key}_encoded_layer'] = LinearOutputBlock(
in_keys=f'{in_key}_encoded_feat',
in_shapes=perception_dict[f'{in_key}_encoded_feat'].out_shapes(
),
out_keys=f'{in_key}_encoded_layer',
output_units=8)
concat_in_keys = [
key for key in perception_dict.keys() if '_encoded_layer' in key
]
perception_dict['hidden_out'] = ConcatenationBlock(
in_keys=concat_in_keys,
in_shapes=sum(
[perception_dict[key].out_shapes() for key in concat_in_keys],
[]),
out_keys='hidden_out',
concat_dim=-1)
self.perception_dict = perception_dict
def __init__(self, obs_shapes: Dict[str, Sequence[int]],
hidden_units: List[int], non_lin: nn.Module):
super().__init__()
self.hidden_units = hidden_units
self.non_lin = non_lin
self.perception_dict: Dict[str, PerceptionBlock] = dict()
# first, flatten all observations
flat_keys = []
for obs, shape in obs_shapes.items():
out_key = f'{obs}_flat'
flat_keys.append(out_key)
self.perception_dict[out_key] = FlattenBlock(
in_keys=obs,
out_keys=out_key,
in_shapes=shape,
num_flatten_dims=len(shape))
# next, concatenate flat observations
in_shapes = [
self.perception_dict[k].out_shapes()[0] for k in flat_keys
]
self.perception_dict["concat"] = ConcatenationBlock(
in_keys=flat_keys,
out_keys='concat',
in_shapes=in_shapes,
concat_dim=-1)
# build perception part
self.perception_dict["latent"] = DenseBlock(
in_keys="concat",
out_keys="latent",
in_shapes=self.perception_dict["concat"].out_shapes(),
hidden_units=self.hidden_units,
non_lin=self.non_lin)
# initialize model weights
module_init = make_module_init_normc(std=1.0)
for key in self.perception_dict.keys():
self.perception_dict[key].apply(module_init)
def build_perception_dict():
""" helper function """
in_dict = build_multi_input_dict(dims=[[100, 1, 16], [100, 1, 8]])
perception_dict = dict()
for in_key, in_tensor in in_dict.items():
# compile network block
net = DenseBlock(in_keys=in_key,
out_keys=f"{in_key}_feat",
in_shapes=[in_tensor.shape[-1:]],
hidden_units=[32, 32],
non_lin=nn.ReLU)
perception_dict[f"{in_key}_feat"] = net
net = ConcatenationBlock(in_keys=list(perception_dict.keys()),
out_keys="concat",
in_shapes=[(32, ), (32, )],
concat_dim=-1)
perception_dict["concat"] = net
return in_dict, perception_dict
def build_perception_dict():
""" helper function """
obs_keys = ["obs_inventory", "obs_screen"]
in_dict = build_multi_input_dict(dims=[[1, 16], [1, 3, 64, 64]],
names=obs_keys)
perception_dict = dict()
# --- block ---
net = DenseBlock(in_keys="obs_inventory", out_keys="obs_inventory_latent",
in_shapes=[in_dict["obs_inventory"].shape[-1:]],
hidden_units=[32, 32],
non_lin=nn.ReLU)
perception_dict["obs_inventory_latent"] = net
# --- block ---
net = VGGConvolutionDenseBlock(in_keys="obs_screen", out_keys="obs_screen_latent",
in_shapes=[in_dict["obs_screen"].shape[-3:]],
hidden_channels=[8, 16, 32],
hidden_units=[32],
non_lin=nn.ReLU)
perception_dict["obs_screen_latent"] = net
# --- block ---
net = ConcatenationBlock(in_keys=list(perception_dict.keys()), out_keys="concat",
in_shapes=[(32,), (32,)], concat_dim=-1)
perception_dict["concat"] = net
# --- block ---
net = LinearOutputBlock(in_keys=["concat"], out_keys="action_move",
in_shapes=[(64,)], output_units=4)
perception_dict["action_move"] = net
# --- block ---
net = LinearOutputBlock(in_keys=["concat"], out_keys="action_use",
in_shapes=[(64,)], output_units=16)
perception_dict["action_use"] = net
return in_dict, perception_dict
def test_mlp_and_concat():
""" perception test """
in_dict = build_multi_input_dict(dims=[[100, 1, 16], [100, 1, 8]])
feat_dict = dict()
for in_key, in_tensor in in_dict.items():
# compile network block
net = DenseBlock(in_keys=in_key,
out_keys=f"{in_key}_feat",
in_shapes=(in_tensor.shape[-1], ),
hidden_units=[32, 32],
non_lin=nn.ReLU)
# update output dictionary
feat_dict.update(net(in_dict))
net = ConcatenationBlock(in_keys=list(feat_dict.keys()),
out_keys="concat",
in_shapes=[(32, ), (32, )],
concat_dim=-1)
out_dict = net(feat_dict)
assert out_dict["concat"].ndim == 3
assert out_dict["concat"].shape[-1] == 64
def __init__(self, obs_shapes: Dict[str, Sequence[int]],
action_logits_shapes: Dict[str, Sequence[int]],
non_lin: Union[str, type(nn.Module)], with_mask: bool):
nn.Module.__init__(self)
self.obs_shapes = obs_shapes
hidden_units = 32
self.perception_dict = OrderedDict()
self.perception_dict['selected_feat'] = DenseBlock(
in_keys='selected_piece',
out_keys='selected_feat',
in_shapes=self.obs_shapes['selected_piece'],
hidden_units=[hidden_units],
non_lin=non_lin)
self.perception_dict['order_feat'] = DenseBlock(
in_keys='ordered_piece',
out_keys='order_feat',
in_shapes=self.obs_shapes['ordered_piece'],
hidden_units=[hidden_units],
non_lin=non_lin)
self.perception_dict['latent'] = ConcatenationBlock(
in_keys=['selected_feat', 'order_feat'],
out_keys='latent',
in_shapes=[[hidden_units], [hidden_units]],
concat_dim=-1)
rotation_out_key = 'cut_rotation_logits' if with_mask else 'cut_rotation'
self.perception_dict[rotation_out_key] = LinearOutputBlock(
in_keys='latent',
out_keys=rotation_out_key,
in_shapes=self.perception_dict['latent'].out_shapes(),
output_units=action_logits_shapes['cut_rotation'][0])
if with_mask:
self.perception_dict['cut_rotation'] = ActionMaskingBlock(
in_keys=['cut_rotation_logits', 'cutting_mask'],
out_keys='cut_rotation',
in_shapes=self.perception_dict['cut_rotation_logits'].
out_shapes() + [self.obs_shapes['cutting_mask']],
num_actors=1,
num_of_actor_actions=None)
self.perception_dict['cut_order'] = LinearOutputBlock(
in_keys='latent',
out_keys='cut_order',
in_shapes=self.perception_dict['latent'].out_shapes(),
output_units=action_logits_shapes['cut_order'][0])
in_keys = ['selected_piece', 'ordered_piece']
if with_mask:
in_keys.append('cutting_mask')
self.perception_net = InferenceBlock(
in_keys=in_keys,
out_keys=['cut_rotation', 'cut_order'],
in_shapes=[self.obs_shapes[key] for key in in_keys],
perception_blocks=self.perception_dict)
# initialize model weights
self.perception_net.apply(make_module_init_normc(1.0))
self.perception_dict[rotation_out_key].apply(
make_module_init_normc(0.01))
self.perception_dict['cut_order'].apply(make_module_init_normc(0.01))
def from_observation_space(
self, observation_space: spaces.Dict) -> InferenceBlock:
"""implementation of :class:`~maze.perception.builders.base.BaseModelBuilder` interface
"""
# get a sample observation
sample = observation_space.sample()
# init perception dict
perception_dict = dict()
in_keys = list()
# --- iterate and process observations ---
for obs_key in observation_space.spaces.keys():
if obs_key not in self.observation_modality_mapping:
BColors.print_colored(
f'ConcatModelBuilder: The observation \'{obs_key}\' could not be found in the '
f'model_builder.observation_modality_mapping and wont be considered as an input to the network.',
BColors.WARNING)
continue
in_keys.append(obs_key)
modality = self.observation_modality_mapping[obs_key]
block_type = self.obs_to_block[modality]
# compile network block
params = self.block_params[modality]
net = block_type(in_keys=obs_key,
out_keys=f"{obs_key}_{block_type.__name__}",
in_shapes=sample[obs_key].shape,
**params)
perception_dict[f"{obs_key}_{block_type.__name__}"] = net
# --- merge latent space observations ---
out_key = ConcatModelBuilderKeys.CONCAT
if ConcatModelBuilderKeys.HIDDEN not in self.obs_to_block \
and ConcatModelBuilderKeys.RECURRENCE not in self.obs_to_block:
out_key = ConcatModelBuilderKeys.LATENT
latent_keys = list(perception_dict.keys())
latent_shapes = [
net.out_shapes()[0] for net in perception_dict.values()
]
net = ConcatenationBlock(in_keys=latent_keys,
out_keys=out_key,
in_shapes=latent_shapes,
concat_dim=-1)
perception_dict[out_key] = net
# --- process with presets ---
if ConcatModelBuilderKeys.HIDDEN in self.obs_to_block:
in_key = out_key
out_key = ConcatModelBuilderKeys.HIDDEN
if ConcatModelBuilderKeys.RECURRENCE not in self.obs_to_block:
out_key = ConcatModelBuilderKeys.LATENT
block_type = self.obs_to_block[ConcatModelBuilderKeys.HIDDEN]
net = block_type(
in_keys=in_key,
out_keys=out_key,
in_shapes=perception_dict[in_key].out_shapes(),
**self.block_params[ConcatModelBuilderKeys.HIDDEN])
perception_dict[out_key] = net
if ConcatModelBuilderKeys.RECURRENCE in self.obs_to_block:
in_key = out_key
out_key = ConcatModelBuilderKeys.LATENT
block_type = self.obs_to_block[ConcatModelBuilderKeys.RECURRENCE]
net = block_type(
in_keys=in_key,
out_keys=out_key,
in_shapes=perception_dict[in_key].out_shapes(),
**self.block_params[ConcatModelBuilderKeys.RECURRENCE])
perception_dict[out_key] = net
# compile inference block
in_shapes = [sample[obs_key].shape for obs_key in in_keys]
net = InferenceBlock(in_keys=in_keys,
out_keys=ConcatModelBuilderKeys.LATENT,
in_shapes=in_shapes,
perception_blocks=perception_dict)
return net