Update ant/pwm_onnx/config.json

ant/pwm_onnx/config.json

@@ -1,11 +1,48 @@
 {
-    "name": "ant_pwm_world_model",
-    "env": "Ant-v5",
-    "model_type": "pwm",
-    "components": {
-        "encoder": "encoder.onnx",
-        "dynamics": "transition.onnx",
-        "reward": "reward.onnx"
+  "name": "ant_pwm_world_model",
+  "env": "Ant-v5",
+  "model_type": "pwm",
+
+  "metadata": {
+    "latent_dim": 512
+  },
+
+  "models": [
+    {
+      "name": "encoder",
+      "framework": "onnx",
+      "format": "onnx",
+      "source": {
+        "path": "encoder.onnx"
+      },
+      "signature": {
+        "inputs": ["obs"],
+        "call_mode": "positional"
+      }
     },
-  "latent_dim": 512
-}
+    {
+      "name": "dynamics",
+      "framework": "onnx",
+      "format": "onnx",
+      "source": {
+        "path": "transition.onnx"
+      },
+      "signature": {
+        "inputs": ["z", "a"],
+        "call_mode": "positional"
+      }
+    },
+    {
+      "name": "reward",
+      "framework": "onnx",
+      "format": "onnx",
+      "source": {
+        "path": "reward.onnx"
+      },
+      "signature": {
+        "inputs": ["z", "a"],
+        "call_mode": "positional"
+      }
+    }
+  ]
+}