easier customization of architecture:

all args can now be set freely when model_size is not specified

tdmpc2/common/parser.py

@@ -40,6 +40,7 @@ def parse_cfg(cfg: OmegaConf) -> OmegaConf:
 	cfg.bin_size = (cfg.vmax - cfg.vmin) / (cfg.num_bins-1) # Bin size for discrete regression
 
 	# Model size
+	if cfg.get('model_size', None) is not None:
 		assert cfg.model_size in MODEL_SIZE.keys(), \
 			f'Invalid model size {cfg.model_size}. Must be one of {list(MODEL_SIZE.keys())}'
 		for k, v in MODEL_SIZE[cfg.model_size].items():

tdmpc2/config.yaml

@@ -49,7 +49,7 @@ vmin: -10
 vmax: +10
 
 # architecture
-model_size: 5
+model_size: ???
 num_enc_layers: 2
 enc_dim: 256
 mlp_dim: 512