Skip to content

Commit e551684

Browse files
yichao-liangYichao Liang
andauthored
Add agent simulation learning workflow (#30)
* Refactor _validate_plan_forward to use option model directly Delegate option execution to option_model.get_next_state_and_num_actions instead of duplicating its termination logic (stuck detection, Wait atom-change checks) and directly accessing its simulator. * Unify backtracking refinement search into shared run_backtracking_refinement Extract the duplicated backtracking loop from run_low_level_search (SeSamE) and _refine_sketch (agent bilevel) into a single run_backtracking_refinement function in planning.py. Both callers now delegate to it with their own sample_fn and validate_fn callbacks, eliminating ~80 lines of duplicated loop/backtracking logic. * Simplify _validate_plan_forward to use run_backtracking_refinement Replace 60 lines of manual option-model execution with a call to run_backtracking_refinement using max_tries=[1] and a sample_fn that returns the pre-grounded options. Remove unused Any import. * Refactor _current_observation/_current_state usage in pybullet_env Move the _current_observation assignment into _reset_state so callers don't need to remember the two-step pattern. Clarify the relationship between _current_observation (backing field) and _current_state (typed read accessor) in docstrings and comments. * Add CFG option to load plan sketch from file instead of LLM Adds agent_bilevel_plan_sketch_file setting that, when set to a file path, loads the plan sketch directly from that file, bypassing the foundation model query. Includes test data files and a unit test. * Remove redundant conditions from Place action in boil_plan_sketch * Scale target joint value based on switch_joint_scale in PyBulletBoilEnv * Refactor _terminal in option model to deduplicate wait-termination logic Extract repeated wait-termination check into _check_wait_termination helper and unify the three _terminal branches into a single definition with config checks inside the function body. * Refactor terminal state logging in _OracleOptionModel to simplify condition checks * Format docstring in get_observation method for improved readability * Refactor PyBulletEnv for readability and better naming - Remove dead/commented-out code and stale self-question comments - Add _VIRTUAL_OBJECT_TYPES constant to replace hardcoded type-name skip lists in _set_state and _get_state - Move env-specific _get_robot_state_dict branches to subclass overrides in pybullet_cover and pybullet_blocks - Extract _get_camera_matrices helper to deduplicate render methods - Extract _get_object_state_dict from _get_state for per-object logic - Move create_pybullet_block/sphere to pybullet_helpers/objects.py - Merge _create_task_specific_objects into _set_domain_specific_state - Rename: _reset_state -> _set_state, _reset_custom_env_state -> _set_domain_specific_state, _extract_feature -> _get_domain_specific_feature - Add docstrings explaining where each method is called from * Regroup PyBulletEnv methods by responsibility and update docstring Reorganize methods into labeled sections (Setup, Public API, Core Loop, State Write/Read, Grasp Management, Action Helpers, Rendering, Utilities) so related functions are adjacent. Update module docstring to document the main public API and state synchronization methods. * Refactor PyBulletEnv: extract _domain_specific_step from step() Add _step_base() and _domain_specific_step() to PyBulletEnv base class. step() now calls _step_base (robot control, physics, grasp) then _domain_specific_step (water filling, heating, etc.), gated by _skip_domain_specific_dynamics flag for kinematics-only mode. Migrate all 15 domain envs to override _domain_specific_step() instead of step(). Envs with pre-step logic (coffee, switch, blocks, cover) still override step() for the pre-step part only. * Update PyBulletEnv module docstring for step() refactoring Document the step_base → domain_specific_step → get_observation flow, _skip_domain_specific_dynamics flag, and _domain_specific_step as an optional override. * Add skip_process_dynamics constructor param to PyBulletEnv Replace direct access to private _skip_domain_specific_dynamics attribute with a public constructor parameter, so callers declare kinematics-only mode at creation time instead of mutating internal state after construction. * Extract run_query_sync helper to remove duplicated async-to-sync bridging Both AgentSessionMixin and AgentExplorer had near-identical wrappers that ran session.query() synchronously via nest_asyncio or asyncio.run. Move that logic into a module-level run_query_sync helper in session_manager and have both callers delegate to it. * Refactor main function: extract and modularize setup logic for clarity and maintainability * Rename agent explorer to agent_plan for clearer naming Distinguishes the grounded-plan explorer from upcoming bilevel variants. AgentExplorer -> AgentPlanExplorer, get_name() 'agent' -> 'agent_plan', file moved to agent_plan_explorer.py, and all callers / docstrings / YAML config examples updated accordingly. * Move AgentSessionMixin into agent_sdk package The mixin is pure agent-session plumbing (session creation, lifecycle, explorer factory) and has no approach-specific logic, so it belongs next to session_manager.py, tools.py, and the sandbox managers rather than in approaches/. * Add AgentBilevelExplorer for sim-learning experiments The explorer asks a Claude agent for a plan sketch, refines it against the approach's current (possibly learned) option model, and rolls the refined plan out in the real env. When the mental model disagrees with reality — e.g. the sketch expects JugFilled after a Wait but the mental model's process dynamics can't produce it — the explorer truncates the plan at the deepest unsatisfiable subgoal (inclusive) so the real-env rollout ends exactly where the disagreement occurs, maximising signal per experiment. Key pieces: - predicators/agent_sdk/bilevel_sketch.py: extracted the sketch build / parse / refine helpers from AgentBilevelApproach as module-level functions so both the approach (solve path) and the new explorer (exploration path) can share them. refine_sketch gains truncate_on_subgoal_fail: the on_step_fail callback snapshots the deepest subgoal failure seen during backtracking, and on exhaustion the captured prefix is returned as the experiment plan. - predicators/explorers/agent_bilevel_explorer.py: new explorer. Reads option_model from tool_context (synced by the approach), builds the sketch prompt via bilevel_sketch, runs refine_sketch with check_subgoals=True, check_final_goal=False, truncate_on_subgoal_fail =True, wraps the result in an option_plan_to_policy that converts OptionExecutionFailure into RequestActPolicyFailure so the episode cleanly terminates at the point of real-env divergence. Stashes the sketch subgoals/options on ToolContext for downstream diffing by the learning approach. - predicators/approaches/agent_bilevel_approach.py: shim methods over bilevel_sketch; behaviour unchanged. - predicators/approaches/agent_planner_approach.py: _create_explorer dispatches both "agent_plan" and "agent_bilevel" through the agent factory path and forwards CFG.explorer as the name. - predicators/explorers/__init__.py: factory branch merged for the two agent-session-backed explorers. - predicators/agent_sdk/tools.py: ToolContext gains last_sketch_subgoals / last_sketch_options fields, populated by the explorer and marked TODO for the learning approach to consume. - tests/explorers/test_agent_bilevel_explorer.py: happy-path, fallback, wait-memory-injection, and deepest-subgoal-failure truncation tests. * Add explorer-specific sample budget and experiment-plan logging - 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. * Add sim-learning approach and synthesis tooling 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. * Update experiment configs for sim-learning - 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. * Refactor sim-learning: extract primitives, add GT simulator factory 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__. * Fix formatting, pylint, and mypy issues for CI compliance - 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. * Update test setup to use test tasks for boil environment and refine test description * Refactor combined model in GT simulator * Fix expected-atoms check to support DerivedPredicates 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. * Skip kinematic reset in PyBullet when only non-kinematic state changed 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. * Support offline dataset learning in AgentSimLearningApproach 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. * Log periodic progress during MCMC parameter fitting 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. * Fix mypy and pylint errors for CI compliance 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. * Apply yapf, isort, and docformatter across the codebase * Inline approach configs into parent files in predicatorv3 * Preserve robot joint config across PyBullet state save/restore 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. * Add 'emcee' to the list of install_requires in setup.py * Force PyBullet FK refresh and skip redundant finger snap 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). * Apply yapf/docformatter to satisfy CI autoformat check * Configure predicatorv3 demos for offline-only sim-learning runs 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. * Add jug orientation handling in PyBulletBoilEnv * Revert getLinkState to PyBullet default (no computeForwardKinematics 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. * Add lo/hi bounds to ParamSpec and skip-MCMC support in fit_params 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. * Build boil param specs dynamically from CFG with lo/hi bounds 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. * Apply lo/hi clamping and configurable noise scale to oracle perturbation 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. * Update installation instructions and add macOS setup script for PyBullet * Update PyBullet version to 3.2.7 and simplify macOS setup script * Refactor liquid color update logic and rename related methods for clarity * Add more debug logging for CogMan and option execution flow * Handle PyBullet physics server crashes with env recreation and retry 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. * Fix jug orientation handling in PyBulletBoilEnv by restoring rotation logic * Update installation instructions and dependencies; remove macOS setup script * Remove mara_robosim dependency from setup.py * Fix get_gt_simulator to use env_name instead of normalized name * Add before/after MSE, likelihood, and param-delta logging for parameter fitting * Use SSE loss and wider walker init so MCMC parameter fitting actually 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. * Move GT simulator components onto module-globals contract 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. * Soften boil parameter-dependent gates with sigmoid weights Replaces hard ``dist < threshold`` indicators in the boil rules with sigmoid-smoothed gates of width ``_SOFT_EPS``. Without smoothing, the LM finite-difference Jacobian is ~zero almost everywhere, and the Hessian identifiability diagnostic is uninformative; emcee also gets a non-flat likelihood as a side effect. State-dependent gates (faucet on/off, jug held) stay hard since they don't enter the parameter likelihood. * Add LM warm-start and Hessian identifiability diagnostic Adds fit_map_lm (Levenberg-Marquardt MAP estimate via SciPy TRF) and log_hessian_identifiability (eigendecompose J^T J/sigma^2 + prior precision to flag sloppy parameter directions). Both run as a single LM pass before MCMC; fit_params now centers walkers on theta_map when code_sim_learning_warm_start_with_lm is set, and short-circuits to it directly when num_mcmc_steps == 0. Also adds compute_residuals (per-feature residual vector LM consumes) and log_sse_breakdown (per-(type, feature) SSE so we can see which features dominate the loss). Two CFG flags gate the new behavior: warm_start_with_lm (default True), log_hessian_identifiability (default False). * Infer process-feature scope from base-sim residuals The agent now declares its own PROCESS_FEATURES alongside PROCESS_RULES and PARAM_SPECS, and the loss is scoped to that declaration (instead of every feature on every type). Before synthesis, the approach runs the base sim on each transition and flags (type, feat) pairs whose prediction diverges from the observation on at least min_hits triples; this set is sent to the agent as a starting hint and used as the eval/test scope until the agent overrides it. The base-sim prediction is precomputed once into base_pred_triples so MCMC's inner loop only evaluates the cheap apply_rules step. create_synthesis_tools now takes the precomputed triples plus the inferred hint, drops the live base_env, and reads PROCESS_FEATURES from exec_ns each call (falling back to the hint when undeclared). * Skip MCMC and use LM warm-start in boil agent config LM warm start alone matches the parameter fit for the current boil oracle program; emcee's MAP-of-walkers cannot improve on it in the time budgeted for 500 steps and routinely lands at higher SSE. Setting num_mcmc_steps to 0 and enabling warm_start_with_lm returns the LM theta_map directly. * Apply yapf and docformatter formatting Cleans up line-wrap and docstring drift across the sim-learning branch so the autoformat CI check is satisfied. Bundles the formatting-only changes for cogman, pybullet_boil, and utils that earlier branch commits left behind, plus minor wraps across the new sim-learning code. * Silence mypy on PyBullet client-id attribute access ``BaseEnv`` doesn't declare ``_physics_client_id`` (only PyBullet subclasses do), and ``_recreate_base_env`` reads it best-effort inside a try block. Bind to a local with type:ignore so mypy stops flagging the access without affecting runtime. * Mark unused action arg in sim_fn to satisfy pylint The simulator callback signature must match StepSimulatorFn's (state, action, params) shape even though apply_rules doesn't use the action. Renaming to _action signals intent and silences pylint's unused-argument check. * Use per-component diff in _set_state to eliminate robot jitter Replace the all-or-nothing kinematic-match gate with a per-component diff: robot pose, each object pose, and held-object identity are each compared against the live PyBullet world and only re-written when they actually differ. _robot_matches_state now compares at the joint level (the prior EE-quaternion path hard-coded roll=0, which spuriously mismatched whenever the wrist had any roll and forced a full reset on every simulate() call). reset_state honors caller-provided joint_positions only when they reconstruct the requested EE pose, falling back to IK otherwise. * Reposition recreated cups and plugs in coffee _set_domain_specific_state _remake_cups creates fresh PyBullet bodies that need to be teleported to their state-specified poses; the per-component diff in _set_state now skips objects whose pose already matches PyBullet, so the explicit _reset_single_object calls ensure freshly-recreated bodies land in the right place. Same treatment for plugs when coffee_machine_has_plug. * Look up predicates lazily in option-model _abstract_function The lambda used to capture predicates at __init__ time, which missed predicates invented later (grammar search) and broke subclasses whose _get_current_predicates depends on attributes not yet set during super().__init__(). * Rename 'kinematics-only' to 'base-sim-only' in docs and test names Terminology cleanup to match how skip_process_dynamics is described elsewhere; the env wraps the full base sim, not just kinematics. * Tighten _robot_matches_state atol so set_state hint forces reset The fast-path joint-match check used atol=1e-2, which let a caller's initial_joint_positions hint be silently treated as "already there" when live joints were within 1e-2 of initial — leaving the EE pose up to ~3e-3 off the requested state. State.allclose compares features at 1e-3, so the test then failed reconstruction. Match the State.allclose tolerance. Also pick up trailing yapf reformatting in two approach files. * Fix flaky test_glib_explorer and test_demo_dataset_loading under pytest-split Both tests pass on master and in isolation but fail on shards 6/8 of CI on this branch. The branch's new tests shifted pytest-split's least_duration distribution so existing tests landed in different shards than on master, exposing pre-existing fragility: - test_glib_explorer[Holding]: score_fn returned 0 (not -inf) for non-target goals, so they weren't filtered. With cover's 7-atom dynamic universe and 10 babbles, ~3.5% of seeds sample no Holding goal and the explorer falls through to a Covers goal, leaving the final state without Holding. Bumped glib_num_babbles to 100 and switched the test's score_fn to return -inf for non-target so the explorer never plans toward an off-target predicate. - test_demo_dataset_loading[10-True-oracle-...]: _ensure_cover_demo_ data_exists only checked file existence. test_demo_dataset's max_initial_demos block writes a 3-trajectory dataset under the cover__demo__oracle__7__... name; the [10-...] case then loaded 3 + generated 3 = 6, expected 10. Added a trajectory-count check so the helper regenerates partial files. * Add unit tests for _robot_matches_state atol and pybullet_helpers.objects - test_robot_matches_state_atol_forces_reset_on_small_drift: locks in the 1e-3 atol regression. A ~5e-3 joint drift (within the previous 1e-2 tolerance, outside the new 1e-3) must NOT be treated as "already there" by the fast-path; _set_state must move the robot back to the requested EE pose at State.allclose precision. - tests/pybullet_helpers/test_objects.py (new): coverage for sample_collision_free_2d_positions, used by 3 PyBullet envs but previously without direct tests. Covers no-overlap (circles and rectangles), bounds, reproducibility across seeds, RuntimeError on impossible packing, and ValueError on unknown shape_type. --------- Co-authored-by: Yichao Liang <ycliang@Yichaos-MacBook-Air.local>
1 parent 7a942a4 commit e551684

73 files changed

Lines changed: 5452 additions & 1694 deletions

File tree

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

predicators/approaches/agent_session_mixin.py renamed to predicators/agent_sdk/agent_session_mixin.py

Lines changed: 12 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,8 @@
88
import os
99
from typing import Any, Dict, List, Optional, Set, Union
1010

11-
from predicators.agent_sdk.session_manager import AgentSessionManager
11+
from predicators.agent_sdk.session_manager import AgentSessionManager, \
12+
run_query_sync
1213
from predicators.agent_sdk.tools import ToolContext, create_mcp_tools, \
1314
get_allowed_tool_list
1415
from predicators.explorers import create_explorer
@@ -127,12 +128,18 @@ def _ensure_agent_session(self) -> None:
127128
tools=tools,
128129
)
129130

131+
extra_names = [
132+
getattr(t, "name", "")
133+
for t in self._tool_context.extra_mcp_tools
134+
]
130135
self._agent_session = AgentSessionManager(
131136
system_prompt=self._get_agent_system_prompt(),
132137
mcp_server=mcp_server,
133138
log_dir=self._get_log_dir(),
134139
model_name=CFG.agent_sdk_model_name,
135-
allowed_tools=get_allowed_tool_list(tool_names),
140+
allowed_tools=get_allowed_tool_list(tool_names,
141+
extra_names=extra_names
142+
or None),
136143
)
137144

138145
if self._agent_session_id is not None:
@@ -179,26 +186,18 @@ def _query_agent_sync(self, message: str) -> List[Dict[str, Any]]:
179186
"""Synchronous wrapper for async agent query."""
180187
self._ensure_agent_session()
181188
assert self._agent_session is not None
182-
try:
183-
loop = asyncio.get_event_loop()
184-
if loop.is_running():
185-
import nest_asyncio # type: ignore[import-untyped,import-not-found] # pylint: disable=import-outside-toplevel
186-
nest_asyncio.apply()
187-
return loop.run_until_complete(
188-
self._agent_session.query(message))
189-
return loop.run_until_complete(self._agent_session.query(message))
190-
except RuntimeError:
191-
return asyncio.run(self._agent_session.query(message))
189+
return run_query_sync(self._agent_session, message)
192190

193191
def _create_agent_explorer(
194192
self,
195193
predicates: Set[Predicate],
196194
options: Set[ParameterizedOption],
195+
name: str = "agent_plan",
197196
) -> BaseExplorer:
198197
"""Create an agent explorer with tool_context and agent_session."""
199198
self._ensure_agent_session()
200199
return create_explorer(
201-
"agent",
200+
name,
202201
predicates,
203202
options,
204203
self._types, # type: ignore[attr-defined]

0 commit comments

Comments
 (0)