| 20 | |
| 21 | |
| 22 | class ConvEncoder(nn.Module): |
| 23 | def __init__(self, obs_dim_dict, module_config_dict, time_steps): |
| 24 | super(ConvEncoder, self).__init__() |
| 25 | self.obs_dim_dict = obs_dim_dict |
| 26 | self.module_config_dict = module_config_dict |
| 27 | self.time_steps = time_steps |
| 28 | |
| 29 | self._calculate_dim() |
| 30 | self._build_network_layer(self.module_config_dict.layer_config) |
| 31 | |
| 32 | def _calculate_dim(self): |
| 33 | input_dim = 0 |
| 34 | for each_input in self.module_config_dict["input_dim"]: |
| 35 | if each_input in self.obs_dim_dict: |
| 36 | # atomic observation type |
| 37 | input_dim += self.obs_dim_dict[each_input] |
| 38 | elif isinstance(each_input, (int, float)): |
| 39 | # direct numeric input |
| 40 | input_dim += each_input |
| 41 | else: |
| 42 | current_function_name = inspect.currentframe().f_code.co_name |
| 43 | raise ValueError(f"{current_function_name} - Unknown input type: {each_input}") |
| 44 | |
| 45 | self.input_dim = input_dim |
| 46 | self.output_dim = self.module_config_dict["output_dim"] |
| 47 | self.hidden_dim = self.module_config_dict["hidden_dim"] |
| 48 | |
| 49 | def _build_network_layer(self, layer_config): |
| 50 | layer_config = self._build_layer_config(layer_config, self.time_steps) |
| 51 | |
| 52 | self.encoder = nn.Sequential(nn.Linear(self.input_dim, self.hidden_dim), nn.ReLU()) |
| 53 | if layer_config["type"] == "Conv1d": |
| 54 | self._build_conv_layer(layer_config) |
| 55 | else: |
| 56 | raise NotImplementedError(f"Unsupported layer type: {layer_config['type']}") |
| 57 | |
| 58 | self.output_layer = nn.Linear(layer_config["out_channels"][-1] * 3, self.output_dim) ## dead |
| 59 | |
| 60 | def _build_layer_config(self, base_config, tsteps): |
| 61 | if tsteps == 5: |
| 62 | out_channels = [20, 10] |
| 63 | kernel_sizes = [2, 2] |
| 64 | strides = [1, 1] |
| 65 | |
| 66 | elif tsteps == 10: |
| 67 | out_channels = [20, 10] |
| 68 | kernel_sizes = [4, 2] |
| 69 | strides = [2, 1] |
| 70 | |
| 71 | elif tsteps == 20: |
| 72 | out_channels = [40, 20] |
| 73 | kernel_sizes = [6, 4] |
| 74 | strides = [2, 2] |
| 75 | else: |
| 76 | raise ValueError(f"Unsupported time_steps for now: {tsteps}") |
| 77 | |
| 78 | return dict( |
| 79 | type=base_config["type"], |