Agent bilevel: LLM-proposed params, domino domain, and agent-SDK tooling overhaul#40
Merged
Conversation
- New setting agent_bilevel_explorer_max_samples_per_step (default 50), separate from the solve-path budget, so the explorer's backtracking cost is independently tunable. - Log the actual experiment plan (option names, objects, params) after refinement so the explorer's output is visible alongside the existing sketch/truncation log lines. - Test config updated to set both budgets explicitly.
AgentSimLearningApproach extends AgentBilevelApproach to learn process dynamics online. Each cycle: the agent synthesizes parameterized process rules via Claude (using run_python / evaluate_simulator / test_simulator MCP tools), parameters are fitted via emcee MCMC, and the learned dynamics are composed with a kinematics-only PyBullet oracle into a combined option model for plan refinement. Key pieces: - predicators/approaches/agent_sim_learning_approach.py: the approach. Initialises with a kinematics-only option model (so AgentBilevelExplorer sees disagreements at process-dynamic subgoals like JugFilled/Boiled), and replaces it with the kin+learned model after each successful synthesis cycle. - predicators/agent_sdk/tools.py: create_synthesis_tools() builds the three MCP tools the synthesis agent uses; extra_mcp_tools field and get_allowed_tool_list(extra_names=) plumbing lets the approach inject them into the session. - predicators/code_sim_learning/: ParamSpec, fit_params (emcee MCMC), compute_mse, LearnedSimulator. - predicators/ground_truth_models/boil/gt_simulator.py: ground-truth process-dynamics simulator for the boil environment. - tests/: approach and param-fitting tests.
- agents.yaml: comment out agent_bilevel preset, add agent_sim_learning with explorer=agent_bilevel and skip_test_until_last_ite_or_early_stopping. - common.yaml: disable failure/test video recording, set num_online_learning_cycles=1 for faster iteration.
Simulation primitives (code_sim_learning/utils.py): - apply_rules(state, rules, params) → ProcessUpdate - merge_updates(base_state, updates, process_features) → State - simulate_step(state, action, base_env, rules, params, features) → State These replace _build_fitted_step_fn, merge_process_updates, _sim_fn_from_rules, and the body of _build_combined_simulator. GT simulator factory (ground_truth_models): - GroundTruthSimulatorFactory ABC + get_gt_simulator(env_name) discovery, following the existing get_gt_options / get_gt_nsrts pattern. - PyBulletBoilGroundTruthSimulatorFactory registered in boil/. - Replaces the hardcoded _load_oracle_simulator in the approach. Oracle ablation flags (settings.py): - agent_sim_learn_oracle_sim_program: load GT rules, skip synthesis. - agent_sim_learn_oracle_sim_params: use GT param values, skip MCMC. Also: kin_env → base_env rename throughout, redundant self._types assignment removed, process_features computed once in __init__.
- yapf + isort autoformatting applied to all touched files. - pylint: fix logging-not-lazy in agent_bilevel_explorer, add broad-except and reimported disables in agent_sim_learning_approach. - mypy: fix base/env variable name collision, add type: ignore on lambda inference, add return type annotations to GT factory methods.
Use utils.abstract to evaluate expected atoms in low-level search so that DerivedPredicates — which require a Set[GroundAtom] rather than a State — are handled correctly alongside regular predicates.
When sequential simulate calls differ only in process features (as in the combined kinematic+learned simulator), reapplying joint positions and tearing down/recreating grasp constraints causes visible arm jitter. Compare robot poses first and skip the kinematic reset path when they already match.
Factor simulator synthesis into a shared _learn_simulator helper so that both learn_from_offline_dataset and learn_from_interaction_results can trigger it on their respective trajectory sources. Also create a separate headless env for parameter fitting so MCMC's thousands of _set_state calls don't thrash the GUI env during training.
Replace the silent run_mcmc call with a manual sample loop that logs step count and best log-probability roughly five times per run, and flushes handlers so the updates appear promptly under buffered logging.
Type-annotate **kwargs on PyBullet env __init__ overrides so mypy doesn't flag them. Initialize attrs used by _domain_specific_step in __init__ (pybullet_coffee, pybullet_switch) to silence defined-outside-init. Type-ignore the emcee import. Fix encoding, unused, protected-access, and redefined-outer-name warnings in the sim-learning tests and agent-SDK tooling.
When a held object's grasp constraint is recreated via _set_state, the gripper frame must match the original world pose exactly — otherwise the recorded base_link->object offset is rotated and the object lands at the wrong world position when the gripper next moves. The State representation only carries (x, y, z, tilt, wrist), so IK during reset can pick a different wrist-roll solution and corrupt the constraint. Thread joint_positions from PyBulletState.simulator_state through reset_state so we skip IK and restore the exact arm configuration. Falls back to IK when joints aren't available (plain State). Also wire wait-termination so refinement and execution can stop Wait when expected atoms hold instead of running to max_num_steps_option_rollout: set _abstract_function on the option model in BilevelPlanningApproach (mirrors AgentPlannerApproach), pass abstract_function into option_plan_to_policy in BilevelProcessPlanningApproach, and inject wait_target_atoms per sample in run_low_level_search.
After resetJointState, PyBullet's getLinkState returns a stale link pose from the previous FK cache, producing 50-500μm drift in the EE pose readback. Pass computeForwardKinematics=1 so world poses are recomputed from current joints on every call. Also skip the explicit finger reset in reset_state when joint_positions are provided: arm_joints already includes the finger joints, so set_joints has restored them to their exact continuous values, and the subsequent loop was overwriting them with the discrete-snapped value from _fingers_state_to_joint. The finger reset still runs on the IK path where set_joints leaves fingers untouched. Together these eliminate the "Could not reconstruct state exactly in reset" warning noise (24 -> 0 on the boil-oracle run).
common.yaml: switch to one demonstration per task with no online learning cycle so launch_simp.py exercises only the offline pipeline. agents.yaml (agent_sim_learning): turn on oracle_sim_program with oracle_sim_params disabled so synthesis fits parameters but starts from the ground-truth program structure.
…flag) Investigation found no measurable difference in reported Cartesian world position or orientation whether the flag is True or False, so the override introduced earlier was not needed.
ParamSpec gains optional lo/hi fields for clamping sampled values. fit_params now reads num_steps from CFG.code_sim_learning_num_mcmc_steps by default; passing 0 (or setting the flag to 0) skips emcee entirely and returns the initial parameter values as the fit result. burn_in is also clamped to num_steps-1 to avoid emcee errors on very short runs. Adds a test covering the skip-MCMC path via CFG.
Replace the module-level BOIL_PARAM_SPECS list with _build_param_specs() so water_fill_speed is derived from CFG.boil_water_fill_speed at call time rather than import time. All specs now carry lo=0.0 to prevent MCMC from sampling physically invalid negative values. get_param_specs() is updated to call _build_param_specs() so per-run CFG values are always reflected.
Oracle parameter perturbation now uses the relative scale from CFG.agent_sim_learn_oracle_sim_param_noise_scale (default 0.2) instead of a hard-coded 20 % figure, and clamps perturbed values to the lo/hi bounds declared in each ParamSpec. Also improves the log message when MCMC is skipped (num_mcmc_steps == 0) so it is clear no fitting occurred.
Converts _build_combined_simulator to an instance method so it can capture self, recreate the base env on pybullet.error, and retry once. Also catches pybullet.error in the oracle option model alongside OptionExecutionFailure. Updates agents.yaml config for testing.
Reproduces the exact domino test tasks from a run (same seed, test_env_seed_offset, and domino flags) and saves a PNG of each test task's initial state, labeling solved vs failed tasks.
Extract resolve_refine_timeout and refine_and_validate_report into bilevel_sketch as the shared refinement + forward-validation + report core. Synthesis (run_refinement_for_synthesis) and the new planner refine_plan_sketch tool both call it, differing only in setup glue: synthesis fits PARAM_SPECS and rebuilds the option model per call, while the planner uses the prebuilt ctx.option_model. Wire refine_plan_sketch into the planner's solve tools when a simulator is available.
…ve prompt
At the start of each _solve, render the task's initial state to
test_images/{taskNNN_}initial_state.png so the agent sees the scene
layout before planning. The prompt now includes a '## Initial State
Image' section pointing to the file when available.
Handles both PyBullet envs (_set_state + render()) and general envs
(render_state) with graceful fallback on failure.
After grasping, the held object may start in shallow penetration from grasp settling. Add allow_shallow_held_object_contacts flag to Phase and wire it through make_move_to_phase, PhaseSkill, and BiRRT. When enabled, initial contacts shallower than the configurable pybullet_birrt_shallow_held_contact_margin (-0.003) are excluded from collision checking so the lift can escape without failing. Applied to the LiftSlightly phase of pick skills. Also adds min contact distance to collision log messages for easier debugging.
Replace the fixed-row staging layout with a grid search that uses oriented-rectangle overlap tests to avoid placing movable dominoes on top of start/target blocks. Returns None (triggering retry) when no collision-free slot is found. Adds _placement_collides, _placement_rect, and _rectangles_overlap helpers with a separating-axis overlap test.
Update domino env __main__ test defaults (seed=1, 1 test task, unfinished state). Rename agent config entry for clarity.
The unfinished-state staging loop placed movable dominoes with an overlap-only collision check, which could leave one inside the gripper's swept grasp footprint of the start block or a target -- especially a perpendicular neighbor a few cm away in y. The domino then lands placed but un-pickable: BiRRT finds no collision-free descent for Pick/MoveToGrasp. Add a grasp-clearance check (_grasp_clearance_blocked): reject a staging spot unless the gripper's swept footprint -- an oriented rectangle with half-extents 0.85x domino width along the long axis and 1.45x along the finger/depth axis, measured from the Fetch gripper -- is clear of every other object. Verified across seeds 0-4: previously seed1 t3, seed2 t4 and seed2 t5 each had an un-pickable movable domino; now every movable domino in all 25 tasks is graspable from init, with no generation slowdown.
Debugging/repro tooling for the domino oracle-samplers runs: - reproduce_domino_failures.py: deterministic, LLM-free reproduction of grasp/place BiRRT infeasibility and the Push parser-drop bug. - replay_domino_sketches.py: replay recorded LLM sketches through the real bilevel refinement to reproduce solve-time failures. - render_unsolved_domino_states.py: annotated init-state PNGs for the unsolved tasks. - plan_sketches/domino_repro_s1t0.txt: example sketch for --agent_bilevel_plan_sketch_file.
Keep these predicates in oracle.yaml (test oracle) but drop them for agent runs. Achieved via a deep-merged ENVS.domino override in agents.yaml instead of the shared envs/all.yaml.
Add a per-phase Phase.validate_ik flag and set it for Pick's MoveToGrasp.
When CFG.pybullet_ik_validate is False, unvalidated PyBullet IK can return a
grasp goal config whose EE pose is numerically close but whose gripper finger
slightly penetrates the very domino being grasped (~1-11mm). BiRRT then rejects
the otherwise-reachable grasp ("no collision-free path"), failing the option
mid-plan even though the grasp pose is feasible (validated IK clears it).
_plan_with_simulator now validates the goal IK when the phase requests it,
without globally enabling validation (which slows transport/place/retreat and
introduces Place/Retreat collision + refinement-budget regressions). Replaying
the recorded domino oracle-samplers sketches confirms this clears the mid-plan
Pick/MoveToGrasp failures (e.g. no_demo seed1.t3 4/5 -> 5/5) with no new
regressions, where global ik_validate=True regressed the same seed to 3/5.
Move the domino reproduction/rendering harnesses into a dedicated scripts/domino_debug/ package and update their usage strings and the cross-import in the probes accordingly. No behavior change.
- probe_cascade.py: run a recorded Pick/Place/Push/Wait sketch through the real option model and log each domino's roll, locating where a topple cascade dies (e.g. Toppled subgoal never propagating past the start block). - probe_infront_drift.py: place a domino at the oracle sampler's generator-faithful pose via the real forward sim and compare the settled pose to the InFront tolerance window. - replay_ikval_sweep.py: replay recorded sketches per task with a configurable ik_validate setting and report FIXED/REGRESSION vs the recorded outcome.
Render the initial-state image at the start of _solve and reference it from the plan-sketch prompt so the sketcher can see the scene layout. Extract the inline image-section construction in agent_planner_approach into a reusable _initial_image_section() helper and thread it through build_solve_prompt.
…oach Drop the global max_initial_demos: 1 default from common.yaml and set it per-env instead (domino: 0; others commented alongside their env blocks). Rename the agent approach entry agent_oracle_hybrid_sim_oracle_samplers_demo to agent_oracle_hybrid_sim_oracle_samplers.
Move domino_restricted_push: True out of the shared envs/all.yaml and into oracle.yaml's ENVS.domino.FLAGS. Agent configs now fall back to the codebase default (False), so the restricted Push option is registered as types=[robot, domino] and the LLM's Push(robot, domino_X) sketch line parses instead of being silently dropped (the seed0/task1 no_demo failure). Also rename the agent approach to agent_oracle_hybrid_sim_oracle_samplers_no_demo and drop its redundant approach-level override.
…pts for analyzing turn percentages in domino tasks
The plain AgentPlannerApproach is the open-loop agent that selects continuous parameters itself, so it should not get refine_plan_sketch (backtracking refinement on a param-free sketch). Move that tool out of the base _get_solve_tool_names and add it via an override on AgentBilevelApproach (still gated on agent_planner_use_simulator), where the agent hands continuous refinement to a search procedure. AgentSimLearningApproach inherits it.
Bound the model's deliberation (default 16000 tokens) so a single response's thinking + text stays under the harness output-token cap, preventing the intermittent 'exceeded the output token maximum' (32000) overflow on hard tasks. Wired through the local sandbox, the shared AgentSessionManager, and the docker runner; gated by the new agent_sdk_thinking_budget_tokens setting (0 leaves it unset).
Let the agent propose per-step continuous parameters inside the plan
sketch (`Option(obj:type)[p1, p2] -> {subgoals}`) instead of always
recovering them by backtracking search. Refinement tries the proposed
params first, then falls back to the registered sampler / uniform
backtracking on failure. Gated by agent_bilevel_use_llm_initial_params
(default False keeps the param-free sketch); agent_bilevel_refine_fallback
controls whether the approach still runs its own post-agent refinement or
the agent must deliver a refine_plan_sketch-validated plan.
When the info-seeking explorer is also on, the proposed params are seeded
as the FIRST candidate of the info-seeking pool (rolled forward,
subgoal-checked, scored) so the ensemble-disagreement argmax chooses among
{LLM guess} U sampled draws, rather than the guess short-circuiting the
probe. One-shot per step, counts toward the node's rollout budget.
- bilevel_sketch: SketchStep.initial_params, parse_initial_params, the
sample_fn LLM-params branch + info-seeking seeding, refine_and_validate
_report now returns the grounded plan.
- agent_bilevel_approach: prompt + tool plumbing for proposing params and
for the tool-validated-only delivery path.
- tools / explorer / synthesis_validation: thread parse_continuous_params
and the new return arity through.
- tests: pooling/seed-wins/seed-loses/infeasible-seed + parser cases.
Turn on agent_bilevel_use_llm_initial_params for the agent_po_predicate_invention_al approach, add InFront to domino's excluded_predicates, drop the now-redundant early-stopping / warm-start flags from that block, and refresh the commented-out A/B reference configs (oracle-samplers vs LLM-params hybrid-sim).
The save_low_level_action_images param on evaluate_option_plan was used in only ~2% of calls across the agent_model_based_planning and agent_oracle_hybrid_sim_llm_params_no_demo runs, and in every case the agent never read the generated per-low-level-action frames (it reasoned off the always-saved per-option scene images instead). Each True triggered hundreds of pybullet renders per task for no consumed benefit. Remove the param from the tool schema, its retrieval, and the low-level rendering branch; per-option image saving is unchanged. Drop the agent-facing docs that pointed at the flag and ./test_images_low_level/.
Run the three CI-pinned autoformatters (yapf 0.32.0, docformatter 1.4, isort 5.10.1) over files accumulated on sim-learning that were never autoformatted. Whitespace/import-ordering only, no logic changes.
Make mypy/pylint pass across the branch's accumulated changes: - agent_sdk: coerce per-solve cost to float (Optional[float] + operand fix) in session_manager/local_sandbox; rename a thinking loop-local in docker_agent_runner that shadowed the SDK thinking config. - approaches: make human_option_control._solve accept _allow_replan so it matches the BilevelProcessPlanningApproach signature (fixes [override] + arguments-differ); type:ignore[override] on maple_q's intentional train_or_test signature; resolve agent_bilevel_plan_sketch_file to an absolute path when given one (fixes test_sketch_from_file); wrap long lines and guard _agent_session_id (defined via the session mixin). - skill_factories/base: drop a shadowed get_link_state import, del the unused objects arg, reflow comments. - domino_debug scripts: add type annotations + docstrings, hoist lazy imports, and mark intentional protected-member probes. - tests: arg-type ignores for the classifier stubs; booleaness fixup.
The task-generator rework on this branch changed the seed-0 chain: the start block (domino_0) now sits at yaw=0 with a wider x-footprint, and the chain bends through a 45-degree turn. The scripted Place params were tuned to the old straight chain, so the second place drove the gripper ~2.5mm into domino_0 and BiRRT's Place/Retreat start was in collision. Set the two Place poses to the generator's intended finished-chain poses (domino_1 -> 0.7854 rad, domino_2 -> 1.5708 rad); Pick/Push unchanged. No collision-margin or skill-factory changes -- the overlap was a real penetration, now eliminated. Fixes test_human_option_control_scripted_domino_solves_task.
isort reordered the imports such that the fully-qualified domino_task_generator import exceeded 80 cols (pylint C0301). Alias the parent package as dtg, matching the other domino_debug modules.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Integration of the
sim-learningwork since #38 (Info-seeking active experiments). This batch advances the agent-driven bilevel pipeline and brings the pybullet domino domain to a working state, plus a substantial agent-SDK tooling overhaul and supporting debug/analysis scripts.Highlights
Agent SDK & bilevel sketching
agent_sdk/tools.pyandagent_sdk/bilevel_sketch.py(the tool surface the agent uses to propose, vet, and refine plan sketches).agent_bilevel_use_llm_initial_params). Wired throughagent_bilevel_approach.pyandagent_sim_learning_approach.py.save_low_level_action_imagesoption fromevaluate_option_plan(used in ~2% of calls, never read by the agent) and its docs.Domino domain
pybullet_domino(composed_env.py→env.py, component refactors for grid/domino), reworked predicates/processes and skill factories.domino_task_generator.pyand added GT sampler tests; oracle/LLM-param configs reach the high-20s/25 success range.Approaches / configs
agent_abstraction_learning_approach.py,agent_closed_loop_approach.py, andchain_reward.py.predicatorv3configs (agents.yaml,envs/all.yaml) andsettings.pyflags.Tooling & tests
scripts/domino_debug/analysis suite (turn diversity/counting, failure reproduction, state rendering, ik-val sweeps, cascade/infront-drift probes).tests/agent_sdk/test_bilevel_sketch_samplers.py, domino GT sampler tests, skill-factory integration, agent-SDK tools.Stats
80 files changed, +6476 / −2877, 266 commits since the #38 merge-base.
Checks
The 4 README CI checks (pytest, mypy, pylint, autoformat) run in CI. Local spot-checks on the most recent changes pass; any pre-existing local-vs-CI formatter/mypy version deltas are unaffected by this batch.