Agent-bilevel: horizon-faithful plan capture + explorer fixes (solves plain domino 5/5 across 5 seeds)#41
Closed
yichao-liang wants to merge 277 commits into
Closed
Agent-bilevel: horizon-faithful plan capture + explorer fixes (solves plain domino 5/5 across 5 seeds)#41yichao-liang wants to merge 277 commits into
yichao-liang wants to merge 277 commits into
Conversation
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.
… moves Switch the fitting loss from per-feature MSE to total SSE (drop the /count in compute_sse) so the Gaussian log-likelihood -0.5*SSE/sigma^2 is in its correct iid form. The previous MSE form silently rescaled per-observation noise by sqrt(count), making walker proposals indistinguishable from each other. Pair this with a wider walker initialization (0.5 * prior_sigma instead of 1% of init_value) so the swarm covers the prior support and emcee stretch moves can actually explore.
Unifies oracle and agent-synthesized simulators behind one loader: read_simulator_components pulls PROCESS_RULES, PARAM_SPECS, and PROCESS_FEATURES out of any namespace (module dict for oracle, exec_ns for agent), and get_gt_simulator now returns the triple including features. merge_updates no longer takes process_features since the rule producer owns that scope.
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.
_get_exploration_strategy never set tool_context.current_task, so refine_plan_sketch/evaluate_option_plan/visualize_state (which default to ctx.current_task when task_idx is omitted) tuned and validated the exploration plan against the stale last TEST task left by test _solve. The agent saw a goal referencing objects the explore task lacked (e.g. Toppled(domino_3) for a 3-object task), so parameter search was meaningless and only tuning-free tasks got solved. Set current_task to the explore task and clear the capture path so an exploration plan can't leak into the next test solve.
Add a GlobalSettings.agent_sdk_reasoning_effort knob ("low"/"medium"/
"high"/"max", or ""/"default" to leave unset) and thread it through both
session managers into ClaudeAgentOptions.effort: the local sandbox
validates the value and the docker sandbox forwards it (plus the thinking
budget) to the container runner. Orthogonal to the per-response thinking
budget; lets hard planning tasks trade latency/cost for more deliberate
reasoning.
execute_plan_forward now tracks total_actions, the first goal-reaching step, and the cumulative actions to reach it, and gains a stop_on_failure mode plus a clean_to_goal property. evaluate_option_plan runs with stop_on_failure=True (aborting at the first failing option, like the real closed-loop executor) and only counts the goal as achieved/captured when it is reached cleanly AND within CFG.horizon low-level steps; when the plan reaches the goal but exceeds the horizon it tells the agent to shorten the plan instead of capturing a false positive. Tests cover the new stop-on-failure and horizon-gating behavior.
…n search When the agent submits + simulator-validates a goal-reaching plan via evaluate_option_plan / refine_plan_sketch but ends with a prose summary whose final text doesn't parse into a sketch, the explorer previously dropped that productive solve to the random-options fallback (a wasted exploration attempt that could block train-driven early stopping). Enable the capture path (keyed to the explore task) and, on an empty parse, rebuild the sketch from the captured plan, grafting each option's continuous params onto SketchStep.initial_params so they seed the info-gain parameter search as candidates rather than being replayed verbatim. Mirrors the test solver's preference for the tool-validated capture over the final text; stale capture is cleared at entry so an exploration plan can't leak into a later test solve.
yichao-liang
added
the
solved: domino 5/5 across 5 seeds
Collaborator
Author
|
Superseded by #42, a clean re-cut off current |
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.
✅ Result: full method solves plain domino in 5/5 tasks across 5/5 seeds
The agent-bilevel method (LLM-proposed plan sketch + continuous params, refined and executed closed-loop) now solves the plain domino domain 5/5 tasks on every one of 5 seeds (25/25), with no oracle samplers.
What's in this PR
Net diff vs
masteris small (11 files); it layers the final fidelity and exploration fixes on top of the already-merged #40 work.Plan capture now mirrors the real executor
The forward-simulation capture path could report a goal the real closed-loop rollout never reaches. Two constraints from the real executor are now enforced before a plan is captured:
stop_on_failure=True): abort at the first failing option (0-action collision / not-initiable / env failure) instead of forward-simulating past it.CFG.horizonlow-level steps is reported asNOT EXECUTABLEand not captured, so the agent shortens the plan instead of stopping on a false positive.Explorer fixes
agent_bilevel_explorerrecovers tool-validated capture and seeds the info-gain search.Config / tooling
agent_sdk_reasoning_effort).horizonraised 400 → 500.Tests
tests/agent_sdk/test_bilevel_sketch_samplers.pytests/explorers/test_agent_bilevel_explorer.pyPlus CI/lint/format cleanup across the branch.