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