| 3 | import inspect |
| 4 | |
| 5 | class BaseModule(nn.Module): |
| 6 | def __init__(self, obs_dim_dict, module_config_dict): |
| 7 | super(BaseModule, self).__init__() |
| 8 | self.obs_dim_dict = obs_dim_dict |
| 9 | self.module_config_dict = module_config_dict |
| 10 | |
| 11 | self._calculate_input_dim() |
| 12 | self._calculate_output_dim() |
| 13 | self._build_network_layer(self.module_config_dict.layer_config) |
| 14 | |
| 15 | def _calculate_input_dim(self): |
| 16 | # calculate input dimension based on the input specifications |
| 17 | input_dim = 0 |
| 18 | for each_input in self.module_config_dict['input_dim']: |
| 19 | if each_input in self.obs_dim_dict: |
| 20 | # atomic observation type |
| 21 | input_dim += self.obs_dim_dict[each_input] |
| 22 | elif isinstance(each_input, (int, float)): |
| 23 | # direct numeric input |
| 24 | input_dim += each_input |
| 25 | else: |
| 26 | current_function_name = inspect.currentframe().f_code.co_name |
| 27 | raise ValueError(f"{current_function_name} - Unknown input type: {each_input}") |
| 28 | |
| 29 | self.input_dim = input_dim |
| 30 | |
| 31 | def _calculate_output_dim(self): |
| 32 | output_dim = 0 |
| 33 | for each_output in self.module_config_dict['output_dim']: |
| 34 | if isinstance(each_output, (int, float)): |
| 35 | output_dim += each_output |
| 36 | else: |
| 37 | current_function_name = inspect.currentframe().f_code.co_name |
| 38 | raise ValueError(f"{current_function_name} - Unknown output type: {each_output}") |
| 39 | self.output_dim = output_dim |
| 40 | |
| 41 | def _build_network_layer(self, layer_config): |
| 42 | if layer_config['type'] == 'MLP': |
| 43 | self._build_mlp_layer(layer_config) |
| 44 | else: |
| 45 | raise NotImplementedError(f"Unsupported layer type: {layer_config['type']}") |
| 46 | |
| 47 | def _build_mlp_layer(self, layer_config): |
| 48 | layers = [] |
| 49 | hidden_dims = layer_config['hidden_dims'] |
| 50 | output_dim = self.output_dim |
| 51 | activation = getattr(nn, layer_config['activation'])() |
| 52 | |
| 53 | layers.append(nn.Linear(self.input_dim, hidden_dims[0])) |
| 54 | layers.append(activation) |
| 55 | |
| 56 | for l in range(len(hidden_dims)): |
| 57 | if l == len(hidden_dims) - 1: |
| 58 | layers.append(nn.Linear(hidden_dims[l], output_dim)) |
| 59 | else: |
| 60 | layers.append(nn.Linear(hidden_dims[l], hidden_dims[l + 1])) |
| 61 | layers.append(activation) |
| 62 | |