Implement policy param

artefacts.yaml

@@ -35,8 +35,8 @@ jobs:
     scenarios:
       settings:
         - name: report_based_waypoint_mission_test
-          run: "uv run dataflow --test-waypoint-report"
+          run: "uv run dataflow --test-waypoint-poses --policy stumbling"
 
         - name: pose_based_waypoint_mission_test
-          run: "uv run dataflow --test-waypoint-poses"
+          run: "uv run dataflow --test-waypoint-poses --policy complete"
 

dataflow/src/dataflow/dataflow.py

@@ -3,6 +3,7 @@
 import dora
 from dora.builder import DataflowBuilder
 from typing_extensions import Annotated
+from typing import Optional
 from pathlib import Path
 
 workspace_path = Path(__file__).parent.parent.parent.parent
@@ -12,6 +13,13 @@
 )
 temp_dataflow_path = output_path / "dataflow.yaml"
 
+policies_folder = Path(__file__).resolve().parent.parent.parent.parent / "policies"
+available_policy_folders = (
+    sorted([p.name for p in policies_folder.iterdir() if p.is_dir()])
+    if policies_folder.exists()
+    else []
+)
+
 
 def _exec_dataflow(dataflow: DataflowBuilder, temp_dataflow_path: Path):
     """Build and/or run the dataflow with dora-rs."""
@@ -20,7 +28,7 @@ def _exec_dataflow(dataflow: DataflowBuilder, temp_dataflow_path: Path):
     dora.run(str(temp_dataflow_path))
 
 
-def _create_base_dataflow() -> DataflowBuilder:
+def _create_base_dataflow(policy_path: Path) -> DataflowBuilder:
     dataflow = DataflowBuilder(name="go2-example-dataflow")
 
     output_path.mkdir(parents=True, exist_ok=True)
@@ -32,6 +40,7 @@ def _create_base_dataflow() -> DataflowBuilder:
         args="--scene generated_pyramid --use-auto-pilot",
         env={
             "OMNI_KIT_ACCEPT_EULA": "YES",
+            "GO2_POLICY_PATH": str(policy_path),
         },
     )
     simulation.add_input("pub_status_tick", "dora/timer/millis/200")
@@ -47,6 +56,7 @@ def _create_base_dataflow() -> DataflowBuilder:
     policy_controller = dataflow.add_node(
         id="policy_controller",
         path="policy_controller",
+        env={"GO2_POLICY_PATH": str(policy_path)},
     )
     policy_controller.add_input("observations", "simulation/observations")
     policy_controller.add_input("clock", "simulation/simulation_time")
@@ -59,6 +69,16 @@ def run_dataflow(
     teleop: Annotated[
         bool, typer.Option(help="Use keyboard teleoperation to control the robot")
     ] = False,
+    policy: Annotated[
+        Optional[str],
+        typer.Option(
+            help=(
+                "Policy folder name inside 'policies' or absolute path. "
+                f"Available: {', '.join(available_policy_folders) if available_policy_folders else 'none detected'}. "
+                "Default: GO2_POLICY_PATH env or 'complete'."
+            )
+        ),
+    ] = None,
     test_waypoint_poses: Annotated[
         bool, typer.Option(help="Run the waypoint poses tests")
     ] = False,
@@ -81,12 +101,25 @@ def run_dataflow(
     if test_waypoint_report or test_all:
         tests.append("test_waypoints_report.py")
 
-    if not tests:
-        print("No tests selected to run. Use --help for options.")
-        # TODO: run default dataflow without tester node (and optionally with teleop)
+    if not policy:
+        policy = os.getenv("GO2_POLICY_PATH")
+        if not policy:
+            policy = "complete"
+
+    if policy in available_policy_folders:
+        policy = policies_folder / policy
+
+    resolved_policy_path = Path(policy).expanduser().resolve()
+
+    if not resolved_policy_path.exists():
+        raise typer.BadParameter(
+            f"Policy path '{resolved_policy_path}' does not exist."
+        )
 
     if teleop:
-        dataflow, simulation, policy_controller = _create_base_dataflow()
+        dataflow, simulation, policy_controller = _create_base_dataflow(
+            resolved_policy_path
+        )
 
         teleop_node = dataflow.add_node(
             id="teleop",
@@ -102,7 +135,9 @@ def run_dataflow(
         _exec_dataflow(dataflow, temp_dataflow_path)
 
     for test in tests:
-        dataflow, simulation, policy_controller = _create_base_dataflow()
+        dataflow, simulation, policy_controller = _create_base_dataflow(
+            resolved_policy_path
+        )
 
         # Add waypoint navigation
         navigator = dataflow.add_node(

nodes/policy_controller/policy_controller/main.py

@@ -1,20 +1,30 @@
 """Policy controller node."""
 
+import os
 from pathlib import Path
 
+import msgs
 from dora import Node
 
-import msgs
 from policy_controller.policy import Policy
 
 
 def main():
     """Receive observations and twist commands and output joint commands."""
     node = Node()
 
+    default_policy_path = (
+        Path(__file__).resolve().parent.parent.parent.parent / "policies" / "complete"
+    )
+    # Try to read policy path from environment variable
+    policy_path_str = os.getenv("GO2_POLICY_PATH", str(default_policy_path))
+    policy_path = Path(policy_path_str)
+    if not policy_path.exists():
+        raise FileNotFoundError(f"Policy path not found at {policy_path}")
+
     policy = Policy(
-        model_path=Path(__file__).parent / "policy" / "policy.pt",
-        config_path=Path(__file__).parent / "policy" / "env.yaml",
+        model_path=policy_path / "policy.pt",
+        config_path=policy_path / "env.yaml",
     )
     control_timestep = policy.config.dt * policy.config.decimation
 

nodes/simulation/simulation/follow_camera.py

@@ -18,7 +18,7 @@ def __init__(self, target_prim_path: str = "/World/Go2/Head_lower"):
             orientation=rot_utils.euler_angles_to_quats(
                 np.array([0, 0, 0]), degrees=True
             ),
-            resolution=(1080 // 4, 720 // 4),
+            resolution=(1080 // 2, 720 // 2),
         )
         self.camera.set_focal_length(1.8)  # Same as the default perspective camera
         self.camera_location = np.array([0.0, 5.0, 2.0])

nodes/simulation/simulation/go2_robot.py

@@ -1,12 +1,14 @@
+import os
 from pathlib import Path
 
-import msgs
 import numpy as np
 from isaacsim.core.utils.rotations import quat_to_rot_matrix
 from isaacsim.core.utils.types import ArticulationAction
 from isaacsim.robot.policy.examples.controllers import PolicyController, config_loader
 from scipy.spatial.transform import Rotation
 
+import msgs
+
 from .height_scan import HeightScanGrid
 
 
@@ -39,13 +41,17 @@ def __init__(
 
         super().__init__(name, prim_path, root_path, usd_path, position, orientation)
 
-        # TODO: load policy config from messages
-        policy_path = (
-            Path(__file__).resolve().parent.parent.parent
-            / "policy_controller"
-            / "policy_controller"
-            / "policy"
+        default_policy_path = (
+            Path(__file__).resolve().parent.parent.parent.parent
+            / "policies"
+            / "complete"
         )
+        # Try to read policy path from environment variable
+        policy_path_str = os.getenv("GO2_POLICY_PATH", str(default_policy_path))
+        policy_path = Path(policy_path_str)
+        if not policy_path.exists():
+            raise FileNotFoundError(f"Policy path not found at {policy_path}")
+
         env_path = policy_path / "env.yaml"
         if not env_path.exists():
             raise FileNotFoundError(f"Env config file not found at {env_path}")

nodes/tester/tester/test_waypoints_poses.py

@@ -19,7 +19,7 @@ def test_receives_scene_info_on_startup(node):
 
 
 @pytest.mark.parametrize("difficulty", [0.1, 0.7, 1.1])
-@pytest.mark.clock_timeout(30)
+@pytest.mark.clock_timeout(15)
 def test_completes_waypoint_mission_with_variable_height_steps(node, difficulty: float):
     """Test that the waypoint mission completes successfully.
 
@@ -29,7 +29,7 @@ def test_completes_waypoint_mission_with_variable_height_steps(node, difficulty:
 
 
 @pytest.mark.parametrize("scene", ["rail_blocks", "stone_stairs", "excavator"])
-@pytest.mark.clock_timeout(50)
+@pytest.mark.clock_timeout(30)
 def test_completes_waypoint_mission_in_photo_realistic_env(node, scene: str):
     """Test that the waypoint mission completes successfully."""
     run_waypoint_mission_test(node, scene, difficulty=1.0)