From 1d894898c1772c779b1018c9ae98e8caae238747 Mon Sep 17 00:00:00 2001 From: vincentpierre Date: Tue, 18 Aug 2020 15:52:21 -0700 Subject: [PATCH 1/3] Fix of the test for multi visual input --- ml-agents/mlagents/trainers/policy/policy.py | 5 ----- .../trainers/torch/model_serialization.py | 18 +++++++++--------- 2 files changed, 9 insertions(+), 14 deletions(-) diff --git a/ml-agents/mlagents/trainers/policy/policy.py b/ml-agents/mlagents/trainers/policy/policy.py index ec6756c84b..0e23f10ec8 100644 --- a/ml-agents/mlagents/trainers/policy/policy.py +++ b/ml-agents/mlagents/trainers/policy/policy.py @@ -43,11 +43,6 @@ def __init__( self.vis_obs_size = sum( 1 for shape in behavior_spec.observation_shapes if len(shape) == 3 ) - self.vis_obs_shape = ( - [shape for shape in behavior_spec.observation_shapes if len(shape) == 3][0] - if self.vis_obs_size > 0 - else None - ) self.use_continuous_act = behavior_spec.is_action_continuous() self.num_branches = self.behavior_spec.action_size self.previous_action_dict: Dict[str, np.array] = {} diff --git a/ml-agents/mlagents/trainers/torch/model_serialization.py b/ml-agents/mlagents/trainers/torch/model_serialization.py index 5e1f16cdec..e664166495 100644 --- a/ml-agents/mlagents/trainers/torch/model_serialization.py +++ b/ml-agents/mlagents/trainers/torch/model_serialization.py @@ -13,15 +13,15 @@ def __init__(self, policy): self.policy = policy batch_dim = [1] dummy_vec_obs = [torch.zeros(batch_dim + [self.policy.vec_obs_size])] - dummy_vis_obs = ( - [torch.zeros(batch_dim + list(self.policy.vis_obs_shape))] - if self.policy.vis_obs_size > 0 - else [] - ) + dummy_vis_obs = [ + torch.zeros(batch_dim + list(shape)) + for shape in self.policy.behavior_spec.observation_shapes + if len(shape) == 3 + ] dummy_masks = torch.ones(batch_dim + [sum(self.policy.actor_critic.act_size)]) dummy_memories = torch.zeros(batch_dim + [1] + [self.policy.m_size]) - # Need to pass all posslible inputs since currently keyword arguments is not + # Need to pass all possible inputs since currently keyword arguments is not # supported by torch.nn.export() self.dummy_input = (dummy_vec_obs, dummy_vis_obs, dummy_masks, dummy_memories) @@ -32,9 +32,9 @@ def __init__(self, policy): if self.policy.use_vec_obs: self.input_names.append("vector_observation") self.dynamic_axes.update({"vector_observation": {0: "batch"}}) - if self.policy.use_vis_obs: - self.input_names.append("visual_observation") - self.dynamic_axes.update({"visual_observation": {0: "batch"}}) + for i in range(self.policy.vis_obs_size): + self.input_names.append(f"visual_observation_{i}") + self.dynamic_axes.update({f"visual_observation_{i}": {0: "batch"}}) if not self.policy.use_continuous_act: self.input_names.append("action_masks") self.dynamic_axes.update({"action_masks": {0: "batch"}}) From 48e77c6b470c65e8620caaa86c9265a396e4be67 Mon Sep 17 00:00:00 2001 From: Ervin Teng Date: Tue, 18 Aug 2020 16:05:04 -0700 Subject: [PATCH 2/3] Make reset block submodule --- ml-agents/mlagents/trainers/torch/encoders.py | 17 +++++++---------- 1 file changed, 7 insertions(+), 10 deletions(-) diff --git a/ml-agents/mlagents/trainers/torch/encoders.py b/ml-agents/mlagents/trainers/torch/encoders.py index b598498e1c..9a36256b1f 100644 --- a/ml-agents/mlagents/trainers/torch/encoders.py +++ b/ml-agents/mlagents/trainers/torch/encoders.py @@ -274,29 +274,26 @@ def __init__(self, height, width, initial_channels, final_hidden): super().__init__() n_channels = [16, 32, 32] # channel for each stack n_blocks = 2 # number of residual blocks - self.layers = [] + layers = [] last_channel = initial_channels for _, channel in enumerate(n_channels): - self.layers.append( - nn.Conv2d(last_channel, channel, [3, 3], [1, 1], padding=1) - ) - self.layers.append(nn.MaxPool2d([3, 3], [2, 2])) + layers.append(nn.Conv2d(last_channel, channel, [3, 3], [1, 1], padding=1)) + layers.append(nn.MaxPool2d([3, 3], [2, 2])) height, width = pool_out_shape((height, width), 3) for _ in range(n_blocks): - self.layers.append(ResNetBlock(channel)) + layers.append(ResNetBlock(channel)) last_channel = channel - self.layers.append(Swish()) + layers.append(Swish()) self.dense = linear_layer( n_channels[-1] * height * width, final_hidden, kernel_init=Initialization.KaimingHeNormal, kernel_gain=1.0, ) + self.sequential = nn.Sequential(*self.layers) def forward(self, visual_obs): batch_size = visual_obs.shape[0] - hidden = visual_obs - for layer in self.layers: - hidden = layer(hidden) + hidden = self.sequential(visual_obs) before_out = hidden.view(batch_size, -1) return torch.relu(self.dense(before_out)) From b6391c7e90380886e3621e7528cba0dcec11f29d Mon Sep 17 00:00:00 2001 From: Ervin Teng Date: Tue, 18 Aug 2020 17:48:45 -0700 Subject: [PATCH 3/3] Fix layers reference --- ml-agents/mlagents/trainers/torch/encoders.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ml-agents/mlagents/trainers/torch/encoders.py b/ml-agents/mlagents/trainers/torch/encoders.py index 9a36256b1f..f5c0f93953 100644 --- a/ml-agents/mlagents/trainers/torch/encoders.py +++ b/ml-agents/mlagents/trainers/torch/encoders.py @@ -290,7 +290,7 @@ def __init__(self, height, width, initial_channels, final_hidden): kernel_init=Initialization.KaimingHeNormal, kernel_gain=1.0, ) - self.sequential = nn.Sequential(*self.layers) + self.sequential = nn.Sequential(*layers) def forward(self, visual_obs): batch_size = visual_obs.shape[0]