Use kw for input arguments to QRDQN

dantp-ai · Apr 2, 2024 · 2b6722f · 2b6722f
1 parent 1c79d19
commit 2b6722f
Showing 1 changed file with 8 additions and 1 deletion.
diff --git a/examples/offline/atari_cql.py b/examples/offline/atari_cql.py
@@ -99,7 +99,14 @@ def test_discrete_cql(args: argparse.Namespace = get_args()) -> None:
     np.random.seed(args.seed)
     torch.manual_seed(args.seed)
     # model
-    net = QRDQN(c, h, w, args.action_shape, args.num_quantiles, args.device)
+    net = QRDQN(
+        c=c,
+        h=h,
+        w=w,
+        action_shape=args.action_shape,
+        num_quantiles=args.num_quantiles,
+        device=args.device,
+    )
     optim = torch.optim.Adam(net.parameters(), lr=args.lr)
     # define policy
     policy: DiscreteCQLPolicy = DiscreteCQLPolicy(