Some encoder details + Add bilevel planning expert

.github/workflows/ci.yml

@@ -67,7 +67,7 @@ jobs:
     - name: Install LAPACK and BLAS
       run: |
         sudo apt-get update
-        sudo apt-get install -y liblapack-dev libblas-dev
+        sudo apt-get install -y liblapack-dev libblas-dev libosmesa6-dev
     - uses: actions/checkout@v4
     - name: Install uv
       uses: astral-sh/setup-uv@v3
@@ -79,5 +79,8 @@ jobs:
     - name: Install the project
       run: uv sync --all-extras --dev
     - name: Pytest
+      env:
+          MUJOCO_GL: osmesa
+          PYOPENGL_PLATFORM: osmesa
       run: |
        uv run pytest tests/

experiments/conf/approach/lpp.yaml

@@ -25,7 +25,7 @@ normalize_plp_actions: true
 permissive_filter_enabled: false
 permissive_filter_max_avg_frac: 0.2
 permissive_filter_max_avg_count: 2
-collision_feedback_enabled: true
+collision_feedback_enabled: false
 collision_feedback_enc: "enc_2"
 collision_template_feedback: true #false: means the feedback is just features/no templates
 collision_feedback_num_buckets: 2
@@ -34,42 +34,44 @@ collision_feedback_max_rounds: 5
 collision_feedback_target_collisions: 0
 prior_version: "uniform" # v1 | v2 | uniform
 prior_beta: 2
-alpha: 0
+alpha: 0.1
 w_clauses: 0.1
 w_literals: 0.05
 w_max_clause: 0.1
 w_depth: 0.05
 w_ops: 0.03
-val_frac: 0.2
+val_frac: 0.0
 val_size: null
 split_seed: 0
 split_strategy: "random"
 program_generation:
   strategy: "py_feature_gen"  # Options: "fixed_grid_v1" "py_feature_gen", "dsl_generator", "feature_generator", offline_loader" "grid_v1_ablated"
   llm_model: "gpt-4.1"
-  demos_included: [0,2,6,8,9]
-  py_feature_gen_prompt: "src/programmatic_policy_learning/dsl/llm_primitives/prompts/py_feature_gen/continuous_example.txt"  # continuous_examplenew_hybrid.txt" new_hybrid
+  demos_included: [0,1,2,3,4,5,6,7,8,9,10] #,6,8, 9
+  skip_rate: 1
+  py_feature_gen_prompt: "src/programmatic_policy_learning/dsl/llm_primitives/prompts/py_feature_gen/continuous_example.txt"  # continuous_example - new_hybrid.txt" new_hybrid
   py_feature_gen_batch_prompt: "src/programmatic_policy_learning/dsl/llm_primitives/prompts/py_feature_gen/batch_prompt_v0.txt"
   num_batches: 1 #keep this 1 for now
   removed_primitive: "" 
   num_features: 50
-  encoding_method: "enc_4" #  "enc_1", "enc_4", "enc_5", "enc_6" are the common options
+  encoding_method: "enc_5" #  "enc_1", "enc_4", "enc_5", "enc_6" are the common options
   loading:
     offline: 0 # 0 for using LLM / 1 for using the following path data
     offline_json_path: "features_rfts_enc5_hybrid_4.1_collision_multi_4.1_trial2.json" #features_stf_demo_only_4.1 features_stf_hint_only_4.1 expanded_stf_notemp_5.2
   negative_sampling:
     enabled: true
     discrete:
-      K: 60
+      K: 16
       local_radius: 2
       w_local: 0.55
       w_struct: 0.3
       w_random: 0.15
     continuous:
+      mode: "manual"  # "mixture" or "manual"
       K: 10
-      local_noise_scale: 0.2
-      lambda_local: 0.7
-      lambda_uniform: 0.3
+      local_noise_scale: 0.06
+      lambda_local: 0
+      lambda_uniform: 1.0
       lambda_traj: 0.0
   # NOTE: The fields below are legacy from earlier experiments.
   # They are not used by the current config unless you switch strategy.

experiments/conf/env/kinder_motion2d.yaml

@@ -3,15 +3,22 @@ num_passages: 0
 make_kwargs:
   base_name: "Motion2D-p${..num_passages}"
   id: "kinder/Motion2D-p${..num_passages}-v0"
-  render_mode: null
+  render_mode: "null"
 
 provider: kinder
 
 description: "KinDER Motion2D-p${.num_passages}. Navigate robot to target region through ${.num_passages} narrow passage(s)."
 
 expert:
-  # _target_: "programmatic_policy_learning.experts.motion2d_expert.Motion2DExpert"
-  _target_: programmatic_policy_learning.approaches.motion2d_expert_approach.Motion2DExpertApproach
+#programmatic_policy_learning.approaches.motion2d_expert_approach.Motion2DExpertApproach
+#motion2d_bilevel_expert_approach.Motion2DBilevelExpertApproach
+  _target_: programmatic_policy_learning.approaches.motion2d_bilevel_expert_approach.Motion2DBilevelExpertApproach
+  # num_passages: ${..num_passages}
+  # max_abstract_plans: 10
+  # samples_per_step: 3
+  # max_skill_horizon: 100
+  # heuristic_name: "hff"
+  # planning_timeout: 60.0
 
 observation_mode: continuous
 action_mode: continuous

experiments/conf/expert/motion2d_expert.yaml

@@ -1,2 +1,8 @@
-# Expert policy for Motion2D: f_1 ∧ f_2 rejection sampling through passages.
-_target_: programmatic_policy_learning.approaches.motion2d_expert_approach.Motion2DExpertApproach
+# Expert policy for Motion2D using bilevel planning.
+_target_: programmatic_policy_learning.approaches.motion2d_bilevel_expert_approach.Motion2DBilevelExpertApproach
+num_passages: 0
+max_abstract_plans: 10
+samples_per_step: 3
+max_skill_horizon: 100
+heuristic_name: "hff"
+planning_timeout: 60.0

experiments/scripts/run_motion2d_rollout.py

@@ -12,6 +12,7 @@
 import argparse
 import logging
 from pathlib import Path
+from typing import Any, Callable, cast
 
 import gymnasium as gym
 import kinder
@@ -30,6 +31,12 @@
 _ACTION_FIELD_NAMES = ["dx", "dy", "dtheta", "darm", "vac"]
 
 
+def _close_env(env: Any) -> None:
+    close_fn = cast(Callable[[], None] | None, getattr(env, "close", None))
+    if close_fn is not None:
+        close_fn()
+
+
 def log_obs(obs: np.ndarray) -> None:
     """Log a structured summary of every observation field."""
     logging.info("--- Robot (indices 0-8) ---")
@@ -186,7 +193,7 @@ def main(args: argparse.Namespace) -> None:
     log_obs(obs)
 
     frames, total_reward, steps, terminated = run_rollout(env, obs, expert)
-    env.close()  # type: ignore[no-untyped-call]
+    _close_env(env)
 
     if terminated:
         logging.info("Done at step %d!", steps)

pyproject.toml

@@ -17,6 +17,8 @@ dependencies = [
    "tomsgeoms2d@git+https://github.com/Princeton-Robot-Planning-and-Learning/prpl-mono.git@bd61b1170856e33c73629245cfedc109d6a12e54#subdirectory=toms-geoms-2d",
    "relational_structs@git+https://github.com/Princeton-Robot-Planning-and-Learning/prpl-mono.git@bd61b1170856e33c73629245cfedc109d6a12e54#subdirectory=relational-structs",
    "generalization_grid_games@git+https://github.com/zahraabashir/generalization_grid_games.git@2af420ecc08f239ff1c7d255b8a790195a9b2865",
+   "kinder-bilevel-planning@git+https://github.com/Princeton-Robot-Planning-and-Learning/kinder-baselines.git#subdirectory=kinder-bilevel-planning",
+   "kinder-models@git+https://github.com/Princeton-Robot-Planning-and-Learning/kinder-baselines.git#subdirectory=kinder-models",
    "moviepy",
    "pandas",
    "pandas-stubs",
@@ -27,6 +29,7 @@ dependencies = [
     "cloudpickle>=2.0.0",
     "tqdm",
     "types-tqdm",
+    "kindergarden",
 ]
 
 [project.optional-dependencies]
@@ -39,9 +42,7 @@ develop = [
     "pytest-pylint>=0.18.0",
     "pytest>=7.2.2",
 ]
-kinder = [
-   "kinder[all]@git+https://github.com/Princeton-Robot-Planning-and-Learning/prpl-mono.git@bd61b1170856e33c73629245cfedc109d6a12e54#subdirectory=kinder",
-]
+
 
 [tool.setuptools.packages.find]
 where = ["src"]
@@ -86,3 +87,8 @@ module = [
     "kinder.*",
 ]
 ignore_missing_imports = true
+
+[dependency-groups]
+dev = [
+    "isort>=8.0.1",
+]

src/programmatic_policy_learning/approaches/experts/kinder_experts.py

@@ -6,10 +6,13 @@
 
 from __future__ import annotations
 
-from typing import Any, Callable
+from typing import Any
 
 import gymnasium as gym
 
+from programmatic_policy_learning.approaches.experts.motion2d_bilevel_experts import (
+    create_motion2d_bilevel_expert,
+)
 from programmatic_policy_learning.approaches.experts.motion2d_experts import (
     create_motion2d_expert,
 )
@@ -19,13 +22,22 @@
     canonicalize_env_name,
 )
 
-ExpertPolicy = Callable[[Any], Any]
+ExpertPolicy = Any
 
 
 def create_kinder_expert(
     env_name: str,
     action_space: gym.spaces.Box,
     seed: int = 0,
+    *,
+    observation_space: Any | None = None,
+    num_passages: int = 0,
+    expert_kind: str = "bilevel",
+    max_abstract_plans: int = 10,
+    samples_per_step: int = 3,
+    max_skill_horizon: int = 100,
+    heuristic_name: str = "hff",
+    planning_timeout: float = 60.0,
 ) -> ExpertPolicy:
     """Build an expert policy for the given KinDER environment.
 
@@ -44,6 +56,25 @@ def create_kinder_expert(
     """
     env_name = canonicalize_env_name(env_name)
     if env_name == "Motion2D":
+        if expert_kind == "bilevel":
+            if observation_space is None:
+                raise ValueError("Motion2D bilevel expert requires observation_space.")
+            return create_motion2d_bilevel_expert(
+                observation_space,
+                action_space,
+                seed=seed,
+                num_passages=num_passages,
+                max_abstract_plans=max_abstract_plans,
+                samples_per_step=samples_per_step,
+                max_skill_horizon=max_skill_horizon,
+                heuristic_name=heuristic_name,
+                planning_timeout=planning_timeout,
+            )
+        if expert_kind != "rejection":
+            raise ValueError(
+                f"Unknown Motion2D expert kind '{expert_kind}'. "
+                "Expected 'rejection' or 'bilevel'."
+            )
         return create_motion2d_expert(action_space, seed=seed)
     raise KeyError(
         f"No KinDER expert configured for '{env_name}'. "

src/programmatic_policy_learning/approaches/experts/motion2d_bilevel_experts.py

@@ -0,0 +1,94 @@
+"""Bilevel-planning experts for Motion2D environments."""
+
+from __future__ import annotations
+
+from typing import Any
+
+
+def _load_bilevel_components() -> tuple[Any, Any]:
+    """Import bilevel-planning modules lazily with a helpful error."""
+    try:
+        # pylint: disable=import-outside-toplevel
+        from kinder_bilevel_planning.agent import BilevelPlanningAgent
+        from kinder_bilevel_planning.env_models import create_bilevel_planning_models
+    except ImportError as exc:  # pragma: no cover - depends on optional install
+        raise ImportError(
+            "Motion2D bilevel expert requires `kinder_bilevel_planning` and its "
+            "dependencies to be installed."
+        ) from exc
+    return BilevelPlanningAgent, create_bilevel_planning_models
+
+
+class Motion2DBilevelPlanningExpert:
+    """Stateful Motion2D expert backed by KinDER bilevel planning."""
+
+    def __init__(
+        self,
+        observation_space: Any,
+        action_space: Any,
+        seed: int = 0,
+        *,
+        num_passages: int = 0,
+        max_abstract_plans: int = 10,
+        samples_per_step: int = 3,
+        max_skill_horizon: int = 100,
+        heuristic_name: str = "hff",
+        planning_timeout: float = 60.0,
+    ) -> None:
+        (
+            bilevel_agent_cls,
+            create_bilevel_planning_models,
+        ) = _load_bilevel_components()
+        env_models = create_bilevel_planning_models(
+            "motion2d",
+            observation_space,
+            action_space,
+            num_passages=int(num_passages),
+        )
+        self._agent = bilevel_agent_cls(
+            env_models,
+            seed,
+            max_abstract_plans=int(max_abstract_plans),
+            samples_per_step=int(samples_per_step),
+            max_skill_horizon=int(max_skill_horizon),
+            heuristic_name=str(heuristic_name),
+            planning_timeout=float(planning_timeout),
+        )
+
+    def reset(self, obs: Any, info: dict[str, Any]) -> None:
+        """Plan from the current initial state."""
+        self._agent.reset(obs, info)
+
+    def step(self) -> Any:
+        """Return the next action in the current plan."""
+        return self._agent.step()
+
+    def update(self, obs: Any, reward: float, done: bool, info: dict[str, Any]) -> None:
+        """Update the wrapped planning agent with the latest transition."""
+        self._agent.update(obs, reward, done, info)
+
+
+def create_motion2d_bilevel_expert(
+    observation_space: Any,
+    action_space: Any,
+    seed: int = 0,
+    *,
+    num_passages: int = 0,
+    max_abstract_plans: int = 10,
+    samples_per_step: int = 3,
+    max_skill_horizon: int = 100,
+    heuristic_name: str = "hff",
+    planning_timeout: float = 60.0,
+) -> Motion2DBilevelPlanningExpert:
+    """Factory that builds a Motion2D bilevel-planning expert."""
+    return Motion2DBilevelPlanningExpert(
+        observation_space,
+        action_space,
+        seed=seed,
+        num_passages=num_passages,
+        max_abstract_plans=max_abstract_plans,
+        samples_per_step=samples_per_step,
+        max_skill_horizon=max_skill_horizon,
+        heuristic_name=heuristic_name,
+        planning_timeout=planning_timeout,
+    )