[Example] Expose compact GAE cat dimension

ghstack-source-id: 701c4e1
Pull-Request: #3775

Author: vmoens

examples/collectors/isaaclab_rnn_ppo_memory.py

@@ -253,6 +253,12 @@ def parse_args() -> argparse.Namespace:
             "'compact' for the compile-friendly compact shifted path."
         ),
     )
+    parser.add_argument(
+        "--gae-compact-cat-dim",
+        choices=["batch", "time"],
+        default="batch",
+        help="Concatenation dimension used by GAE shifted='compact'.",
+    )
     parser.add_argument(
         "--deactivate-vmap",
         action=argparse.BooleanOptionalAction,
@@ -342,6 +348,7 @@ def main() -> None:
         value_network=full_value,
         average_gae=False,
         shifted=gae_shifted,
+        compact_cat_dim=args.gae_compact_cat_dim,
         deactivate_vmap=args.deactivate_vmap,
         device=train_device,
     )