Add bilevel_models primitive: expose the planner's models to the agent

[merlerm]

Summary

Adds the bilevel_models primitive: it exposes the PRPL bilevel planner's models (predicates, operators, parameterized skills, and the transition simulator) to the LLM coding agent, without giving it the planner itself. The agent composes these pieces into one frozen, generalized program instead of running per-instance search.

This is the code behind the agentic bilevel experiments we went over in today's call — e.g. the obstruction2d models-ON result (agent program 90% vs. the SeSamE planner 33%) and the clutteredstorage2d hard-env analysis.

Stacked on the bilevel-planning per-instance baseline (#114).

What's in it

• BilevelModels primitive (primitives/bilevel_models.py) — env-bound access to abstract state / goal atoms / objects / operators / skills, plus run_skill and the models' transition_fn for simulated previews. Lazy-constructed so it's cheap until used.
• Shared model builder (utils/bilevel.py) — one build_sesame_models used by both the planner baseline and the primitive, plus infer_bilevel_mapping so a plain KinderGeom2DEnv the agent builds by hand still resolves its bilevel family.
• Planner anti-cheat (utils/episode.py) — when bilevel_models is granted, a generated approach.py that references the planner (run_sesame, backtracking refiner, planning graph) is rejected at load. The agent must use the models, not call the search.
• Conditional isolation — the bilevel packages are an optional bilevel extra (lazy imported). A sandbox only mounts kinder-baselines and uv sync --extra bilevel when the primitive is granted; otherwise no bilevel source or packages reach the agent, and the prompt descriptions are stripped. This keeps the task structure the models encode out of runs that shouldn't see it.
• Red-teaming — deterministic + adversarial models-OFF probes assert a models-OFF sandbox has no path to recover the models (import, lazy builder, filesystem search, symbol grep, self-install).

Deferred: black-box support (follow-up PR)

bilevel_models runs in clearbox/local mode only. In black-box mode it is rejected up front (loud error) rather than silently half-working.

Supporting it in black-box is feasible but a standalone piece of work, and carries a design trade-off, so it's intentionally a follow-up:

• The generic host-object proxy refuses dunders, but the agent's programs iterate and str() the symbolic objects (for a in atoms, str(atom)), so the symbolic layer (Atom / Object / Type / Predicate / Operator) needs a real serialization into client-side objects, with continuous states kept as host tokens. Comparable in scope to the crv_motion_planning remote-module path.
• Design note: black-box exists to withhold task structure so we can measure the value of agency; bilevel_models hands over the full symbolic decomposition. Exposing it in black-box relaxes what "black-box" means, so it should be a deliberate choice, not just plumbing.

Testing

• Unit tests for the primitive, shared builder, anti-cheat, env-id inference, and the render-server class-primitive path.
• Sandbox tests for the conditional bilevel mount/extra (Docker + Apptainer) and the models-OFF red-team container probe.
• Regression tests for the three container-path fixes in the final commit (genplan mount keeps primitive_descriptions.py; genplan syncs the bilevel extra when requested; black-box manifest rejects bilevel_models).
• mypy, pylint (10/10), and formatters clean.

🤖 Generated with Claude Code

[Copilot]

Pull request overview

This PR introduces a new bilevel_models primitive that exposes SeSamE bilevel planner models (symbolic operators/predicates/skills + transition simulator) to agent code, while keeping the planner/search itself inaccessible. It also adds conditional sandbox mounting/syncing so bilevel sources/deps are only present when the primitive is granted, plus tests and red-team probes to validate “models OFF” isolation.

Changes:

• Add BilevelModels primitive and a shared lazy model builder (build_sesame_models, infer_bilevel_mapping) used by both the primitive and the bilevel baseline.
• Add anti-cheat load-time rejection of generated programs that reference planner symbols when bilevel_models is present.
• Add conditional Docker/Apptainer sandbox mounting and uv sync --extra bilevel only when bilevel_models is requested, plus unit/integration tests (including models-OFF probes).

Reviewed changes

Copilot reviewed 27 out of 27 changed files in this pull request and generated 4 comments.

Show a summary per file

File	Description
tests/utils/test_episode.py	Adds unit tests for the anti-cheat load-time planner-reference rejection.
tests/utils/test_env_server.py	Ensures black-box primitive manifest rejects bilevel_models loudly.
tests/utils/test_docker_sandbox.py	Expands mount filtering to optionally include bilevel sources; adds models-OFF red-team container probe.
tests/utils/test_bilevel.py	Adds tests for shared bilevel model builder and env-id inference consistency with configs.
tests/utils/test_apptainer_sandbox.py	Validates Apptainer command composition for conditional bilevel mounting/syncing.
tests/primitives/test_bilevel_models.py	Adds primitive-level tests for laziness, symbolic layer access, and run_skill.
tests/mcp/test_local_render.py	Verifies in-sandbox primitive rebuilding supports both env-bound functions and env-bound classes.
tests/approaches/test_llm_genplan_approach.py	Ensures genplan container runs include bilevel extras only when bilevel_models is requested.
tests/approaches/test_bilevel_planning_approach.py	Updates missing-mapping test to an unsupported env case.
src/robocode/utils/episode.py	Adds anti-cheat check rejecting planner references when bilevel_models is present.
src/robocode/utils/docker_sandbox.py	Adds conditional copying/mounting of kinder-baselines and --extra bilevel syncing.
src/robocode/utils/bilevel.py	Introduces shared model construction and env-id→mapping inference.
src/robocode/utils/apptainer_sandbox.py	Adds conditional bilevel bind + extra syncing for Apptainer runs/genplan.
src/robocode/primitives/bilevel_models.py	Introduces the BilevelModels env-bound primitive (symbolic access + run_skill).
src/robocode/primitives/init.py	Registers bilevel_models and re-exports descriptions from host-only module.
src/robocode/primitive_specs.py	Adds bilevel_models metadata, env-dependent binding info, and black-box rejection.
src/robocode/primitive_descriptions.py	Splits human-facing primitive descriptions out of sandbox-safe metadata.
src/robocode/mcp/local_render.py	Updates sandbox primitive builder to handle env-dependent class primitives.
src/robocode/environments/kinder_geom2d_env.py	Adds env-id-based bilevel mapping inference fallback when mapping not provided.
src/robocode/approaches/llm_genplan_approach.py	Switches to host-only primitive descriptions import and forwards include_bilevel to container runner.
src/robocode/approaches/bilevel_planning_approach.py	Reuses shared build_sesame_models instead of duplicating builder logic.
src/robocode/approaches/agentic_cdl_approach.py	Switches to host-only primitive descriptions import.
src/robocode/approaches/agentic_base.py	Switches to host-only primitive descriptions import.
pyproject.toml	Moves bilevel deps into an optional bilevel extra.
integration_tests/red_team_sandbox.py	Adds “models OFF” adversarial prompts and detection logic.
docker/entrypoint.sh	Allows passing optional uv sync extra args via ROBOCODE_UV_EXTRA_ARGS.
docker/Dockerfile	Documents runtime-only bilevel extra install and optional kinder-baselines mount.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[tomsilver]

interesting. can we change this so that run_skill doesn't internally sample, but instead, the continuous parameters need to be passed in? I don't think that sampling once makes a lot of sense

[merlerm]

that's a very good point, I had actually missed this entirely. I want to look more closely at this tomorrow. I guess a way to do it is expose the parameter space as part of the skill, so claude can decide them in whatever way it wants. We can also provide a sampling helper but I wouldn't to start

[tomsilver]

change this "instruction"

[merlerm]

do you mean the "use them to ..." part? or simplifying in general?

[tomsilver]

Yeah I meant the "use them to" part!

[merlerm]

Fixed the easy copilot stuff but I want to have a look at the rest more carefully, today I was busy with emergency reviews for ARR so I will get to this tomorrow hopefully

[merlerm]

I decided to cut out all of the run_skill stuff, I was debating even removing the BilevelModels class in general. I think this way the agent just gets a very thin interface that helps with readability (all modules are in the same place) but it needs to figure out how to use them by itelf. Skills need continuous params, these can be sampled with the sample_parameters function that each controller object exposes, but we don't explicitly say this, rather the agent has to figure it out. I initially had a line about this, but I found that removing it almost lets the agent think about stuff a bit more creatively, instead of hard conditioning on just sampling. We can decide if this is good or not.

A test run on stickbutton easy: the agent spent some time at the beginning exploring the bilevel planning module:

Agent's final report

## Report

1. Predicates for StickButton2D (stickbutton2d.py:69-82)

• Grasped(?robot, ?stick) — robot is holding the stick. Computed at stickbutton2d.py:94-98: uses get_suctioned_objects(x, robot); if stick is in the suctioned set, adds Grasped, else adds HandEmpty.
• HandEmpty(?robot) — negation of the above; same check (:95-98).
• Pressed(?button) — button's RGB color matches config.button_pressed_rgb (:104-111), checked via np.allclose.
• RobotAboveButton(?robot, ?button) — robot's geometry intersects button's geometry (:130-135), via object_to_multibody2d + Geom2D.intersects.
• StickAboveButton(?stick, ?button) — stick's geometry intersects button's geometry (:137-140).
• AboveNoButton() (0-arity) — true iff neither robot nor stick intersects any button (:142-144).
• No NEq/OnTarget predicates exist in this file (that's terminology from elsewhere, not used here).

2. Operators (7 total), all at stickbutton2d.py:166-266:

• RobotPressButtonFromNothing(?robot, ?button): pre {HandEmpty(robot), AboveNoButton()}; add {Pressed(button), RobotAboveButton(robot,button)}; del {AboveNoButton()} (:166-178)
• RobotPressButtonFromButton(?robot, ?button, ?from_button): pre {HandEmpty(robot), RobotAboveButton(robot,from_button)}; add {Pressed(button), RobotAboveButton(robot,button)}; del {RobotAboveButton(robot,from_button)} (:181-193)
• PickStickFromNothing(?robot, ?stick): pre {HandEmpty(robot), AboveNoButton()}; add {Grasped(robot,stick)}; del {HandEmpty(robot)} (:196-207)
• PickStickFromButton(?robot, ?stick, ?from_button): pre {HandEmpty(robot), RobotAboveButton(robot,from_button)}; add {Grasped(robot,stick), AboveNoButton()}; del {HandEmpty(robot), RobotAboveButton(robot,from_button)} (:210-225)
• StickPressButtonFromNothing(?robot, ?stick, ?button): pre {Grasped(robot,stick), AboveNoButton()}; add {StickAboveButton(stick,button), Pressed(button)}; del {AboveNoButton()} (:228-240)
• StickPressButtonFromButton(?robot, ?stick, ?button, ?from_button): pre {Grasped(robot,stick), StickAboveButton(stick,from_button)}; add {StickAboveButton(stick,button), Pressed(button)}; del {StickAboveButton(stick,from_button)} (:243-255)
• PlaceStick(?robot, ?stick): pre {Grasped(robot,stick)}; add {HandEmpty(robot), AboveNoButton()}; del {Grasped(robot,stick)} (:258-266)

Skills bind operators to controllers at :288-305, using create_lifted_controllers imported from kinder_models.kinematic2d.envs.stickbutton2d.parameterized_skills (:21-23).

3. Controllers are closed-loop, plan-executing objects implementing GroundParameterizedController (base class Kinematic2dRobotController in /robocode/third-party/kinder-baselines/kinder-models/src/kinder_models/kinematic2d/utils.py:15-128). Each controller: reset(x, params) clears cached plan (:97-103); step() lazily builds a waypoint plan via _generate_plan/_generate_waypoints on first call, then pops one action per call (:109-114); terminated() returns True once the plan list is empty (:105-107). So it's a stateful closed-loop object returning one action per step(), not a single action nor a raw generator.

• Pick skill (GroundPickStickController, .../parameterized_skills.py:23-161): sample_parameters picks a grasp ratio and arm length (:35-43); _calculate_grasp_point (:45-85) and _calculate_robot_position (:87-124) do geometry; _generate_waypoints (:126-158) returns a 6-waypoint SE2 path (retract arm, move to safe height, align x, orient, descend, extend arm) with vacuum on (_get_vacuum_actions returns (0.0, 1.0), :160-161).
• Press skill (GroundRobotPressButtonController, :200-237): _generate_waypoints (:218-234) returns a single waypoint moving the robot base to intersect the button (target_x, target_y = button_x, button_y), vacuum off (:236-237). GroundStickPressButtonController (:240-286) computes dx, dy from the stick's far end to the button and moves the robot by that offset while keeping vacuum on (holding stick), :270-286.

4. Goal: goal_deriver at stickbutton2d.py:150-157 — "The goal is to press all buttons": builds Object(f"button{i}", CircleType) for i in range(num_buttons) and adds GroundAtom(Pressed, [button]) for each, wrapped in RelationalAbstractGoal(atoms, state_abstractor). No function literally named get_goal_atoms exists anywhere under kinder-baselines (grep returned nothing); the equivalent is goal_deriver.

5. State abstraction / reachability: No function named get_abstract_state exists (grep empty); the equivalent is state_abstractor (stickbutton2d.py:85-147). There's no explicit "stick needed vs directly reachable" predicate computed geometrically ahead of time — reachability is implicit in whether RobotAboveButton/StickAboveButton currently hold, and the planner (SeSamE/PDDL search) decides whether to use RobotPressButtonFromNothing/FromButton or the Stick... operators based on precondition satisfiability, not a precomputed "needs stick" predicate. The only geometric check is intersection-based, at :130-140:

if any(robot_geom_sub.intersects(button_geom) for robot_geom_sub in robot_geom):
if any(robot_geom_sub.intersects(button_geom) for robot_geom_sub in robot_geom):
    atoms.add(GroundAtom(RobotAboveButton, [robot, button]))
...
if stick_geom.intersects(button_geom):
    atoms.add(GroundAtom(StickAboveButton, [stick, button]))

using object_to_multibody2d (imported from kinder.envs.kinematic2d.utils, :114-117,126-128) and Geom2D.intersects (tomsgeoms2d.structs, imported :35).

6. Building the bilevel_models primitive bundle:

• Direct construction API: from kinder_bilevel_planning.env_models.kinematic2d.stickbutton2d import create_bilevel_planning_models then call create_bilevel_planning_models(observation_space, action_space, num_buttons=N) (signature at stickbutton2d.py:38-40).
• Generic dynamic loader (preferred, used by robocode): from kinder_bilevel_planning.env_models import create_bilevel_planning_models (.../env_models/__init__.py:13-15), called as create_bilevel_planning_models("stickbutton2d", observation_space, action_space, num_buttons=N) — it dynamically loads kinematic2d/stickbutton2d.py by name (__init__.py:20-53).
• The actual robocode harness that builds this is /robocode/src/robocode/utils/bilevel.py:48-71, function build_sesame_models(env):

from kinder_bilevel_planning.env_models import create_bilevel_planning_models
assert env.bilevel_env_name is not None, ...
return create_bilevel_planning_models(
    env.bilevel_env_name,
    env.observation_space,
    env.action_space,
    **env.bilevel_env_model_kwargs,
)

it found that Pick has sample_parameters for example, just decided to not use it in the later approach, hardcoding some parameters instead.

This was the final approach:

Final approach

"""Approach for kinder/StickButton2D-b1-v0 using bilevel planning models.

Strategy (hand-composed sequencing of the bilevel_models skills/controllers,
not the SeSamE planner):

- There is exactly one button. It sits either on the "floor" (directly
  touchable by the robot base) or on the "table" (a raised region the robot
  base cannot enter, but the stick -- exactly as long as the table is deep --
  can be used to touch it).
- Floor button: run the `RobotPressButtonFromNothing` ground controller
  (from bilevel_models), which drives the robot base onto the button.
- Table button: ground and run the `PickStickFromNothing` controller to
  grasp the stick, then perform our own translation of the robot (still
  holding the stick, vacuum on) so the button falls anywhere within the
  stick's rectangle span -- NOT necessarily at its far corner. (The canned
  `StickPressButtonFromNothing` controller insists on moving the stick's far
  corner exactly onto the button, which for buttons on the near half of the
  table requires the robot base to leave the world; since the stick is as
  long as the table is deep, it is almost always already positioned to cover
  the button in y, so only a small, safe adjustment plus an x-alignment is
  actually needed.)
"""

from __future__ import annotations

from typing import Any

import numpy as np

# Fixed StickButton2D config constants (world/table geometry is never
# randomized across episodes -- only robot/stick/button poses are).
_WORLD_MIN_X = 0.0
_WORLD_MAX_X = 3.5
_WORLD_MIN_Y = 0.0
_WORLD_MAX_Y = 2.5
_TABLE_CENTER_Y = _WORLD_MIN_Y + (_WORLD_MAX_Y - _WORLD_MIN_Y) / 2
_TABLE_HEIGHT = (_WORLD_MAX_Y - _WORLD_MIN_Y) / 2
_TABLE_Y_MIN = _TABLE_CENTER_Y - _TABLE_HEIGHT / 2
_TABLE_Y_MAX = _TABLE_CENTER_Y + _TABLE_HEIGHT / 2


class GeneratedApproach:
    def __init__(self, action_space: Any, observation_space: Any, primitives: dict) -> None:
        self.action_space = action_space
        self.observation_space = observation_space
        self.bm = primitives["bilevel_models"]
        self._skill_by_name = {
            skill.operator.name: skill for skill in self.bm.models.skills
        }
        self._max_dx = float(action_space.high[0])
        self._max_dy = float(action_space.high[1])

        self._phase = "done"  # "press_floor" | "pick" | "align" | "done"
        self._controller = None
        self._action_queue: list[np.ndarray] = []
        self._zero_action = np.zeros(self.action_space.shape, dtype=np.float32)

    def reset(self, state: np.ndarray, info: dict) -> None:
        obs = np.asarray(state, dtype=np.float32)
        robot = self.bm.get_objects(obs, "crv_robot")[0]
        stick = self.bm.get_objects(obs, "rectangle")[0]
        button = self.bm.get_objects(obs, "circle")[0]
        self._robot, self._stick, self._button = robot, stick, button

        state_obj = self.bm.models.observation_to_state(obs)
        button_y = state_obj.get(button, "y")

        self._action_queue = []
        # The table spans the entire x-width of the world, so the robot
        # (which always starts below the table) can never cross it to reach
        # the floor region above the table either. Any button at or above
        # the table's near edge therefore requires the stick; only buttons
        # strictly below the table are directly reachable.
        if button_y >= _TABLE_Y_MIN:
            # Table/upper-floor button: grasp the stick first.
            self._phase = "pick"
            skill = self._skill_by_name["PickStickFromNothing"]
            ground = skill.controller.ground([robot, stick])
            ground.reset(state_obj, (0.5, 0.15))
            self._controller = ground
        else:
            # Floor button: press directly with the robot base.
            self._phase = "press_floor"
            skill = self._skill_by_name["RobotPressButtonFromNothing"]
            ground = skill.controller.ground([robot, button])
            ground.reset(state_obj, 0.0)
            self._controller = ground

    def _translate_actions(self, dx: float, dy: float, vac: float) -> list[np.ndarray]:
        num_steps = int(
            max(1, np.ceil(abs(dx) / self._max_dx), np.ceil(abs(dy) / self._max_dy))
        )
        step_dx = dx / num_steps
        step_dy = dy / num_steps
        action = np.array([step_dx, step_dy, 0.0, 0.0, vac], dtype=np.float32)
        return [action.copy() for _ in range(num_steps)]

    def _start_align_phase(self, obs: np.ndarray) -> None:
        """Compute a translation-only maneuver so the button falls within the
        stick's rectangle span (not necessarily at its far corner), then
        aligns x. Executed as a precomputed (open-loop) action queue."""
        state_obj = self.bm.models.observation_to_state(obs)
        stick_x = state_obj.get(self._stick, "x")
        stick_y = state_obj.get(self._stick, "y")
        width = state_obj.get(self._stick, "width")
        height = state_obj.get(self._stick, "height")
        button_x = state_obj.get(self._button, "x")
        button_y = state_obj.get(self._button, "y")

        margin = 0.02
        if button_y < stick_y + margin:
            dy = (button_y + margin) - stick_y
        elif button_y > stick_y + height - margin:
            dy = (button_y - margin) - (stick_y + height)
        else:
            dy = 0.0
        dx = button_x - (stick_x + width / 2)

        actions = self._translate_actions(dx, dy, vac=1.0)
        # A few settle steps (still holding) so the per-step button-touch
        # check has a chance to trigger once we have arrived.
      actions.append(np.array([0.0, 0.0, 0.0, 0.0, 1.0], dtype=np.float32))
      self._action_queue = actions
      self._phase = "align"
      self._controller = None

  def get_action(self, state: np.ndarray) -> np.ndarray:
      obs = np.asarray(state, dtype=np.float32)

      if self._phase in ("press_floor", "pick"):
          assert self._controller is not None
          if self._controller.terminated():
              if self._phase == "press_floor":
                  self._phase = "done"
                  self._controller = None
              else:  # "pick" finished -> begin alignment maneuver.
                  self._start_align_phase(obs)
          else:
              return np.asarray(self._controller.step(), dtype=np.float32)

      if self._phase == "align":
          if self._action_queue:
              return self._action_queue.pop(0)
          self._phase = "done"

      return self._zero_action.copy()

continuous params in ground.reset were handled by the agent directly. Interestingly, it decided to use a sort of hybrid approach, skipping the skill to press the button with the stick:

The canned StickPressButtonFromNothing controller insists on moving the stick's far corner exactly onto the button, which for buttons on the near half of the table requires the robot base to leave the world... so only a small, safe adjustment plus an x-alignment is actually needed.

So it ended up just using the skill to grasp the stick and wrote its own low level code directly to sweep across the table and press the button. I think this type of behavior is pretty cool.
The only think I am still uncertain about is if we should add some line back in the prompt to point to the samplers directly, e.g.

... their controllers (.ground(objects).controller, parameterized by continuous values that sample_parameters(state, rng) can draw), ...

[merlerm]

Oh and another thing, I saw that at first the agent was just looking into our approaches dir and saw the BilevelPlanningApproach, then spent a lot of time analyzing it and reasoning about how to use the planner just like it did. In general all these approaches are just baselines or different implementations, but the agent at runtime does not need anything beyond the base one, so I stripped the rest

[tomsilver]

Oh and another thing, I saw that at first the agent was just looking into our approaches dir and saw the BilevelPlanningApproach, then spent a lot of time analyzing it and reasoning about how to use the planner just like it did. In general all these approaches are just baselines or different implementations, but the agent at runtime does not need anything beyond the base one, so I stripped the rest

Interesting, is it still possible for the agent to do this? Maybe we should make it so that the baseline approach files are not copied into Docker?

[tomsilver]

The agent's final report mentions "get_goal_atoms" and "get_abstract_state" as functions it looked for but did not exist--did it confuse itself, or do we have those mentioned in prompts somewhere?

[merlerm]

Oh and another thing, I saw that at first the agent was just looking into our approaches dir and saw the BilevelPlanningApproach, then spent a lot of time analyzing it and reasoning about how to use the planner just like it did. In general all these approaches are just baselines or different implementations, but the agent at runtime does not need anything beyond the base one, so I stripped the rest

Interesting, is it still possible for the agent to do this? Maybe we should make it so that the baseline approach files are not copied into Docker?

Commit 8757d00 should be taking care of this

[merlerm]

The agent's final report mentions "get_goal_atoms" and "get_abstract_state" as functions it looked for but did not exist--did it confuse itself, or do we have those mentioned in prompts somewhere?

Good catch! These are from our BilevelModels module, they just wrap the corresponding implementations from kinder and add observation_to_state built in so the agent doesnt waste time figuring out the conversion. In this instance it looks like the agent skipped our primitives code and went directly into the bilevel planning source code, where these are not defined.
In general I guess this is another pointer that this class is not super useful, and we could just expose robocode.utils.bilevel directly and let the agent figure it out... not entirely sure what the best way is