Skip to content

Extract per-skill sampler learning into SamplerLearningMixin #45

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
yichao-liang merged 1 commit into from
Jul 4, 2026
Merged

Extract per-skill sampler learning into SamplerLearningMixin #45

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
315 changes: 11 additions & 304 deletions predicators/approaches/agent_sim_learning_approach.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -30,10 +30,10 @@

from predicators import utils
from predicators.agent_sdk.tools import SAMPLER_SYNTHESIS_TOOL_NAMES, \
SYNTHESIS_TOOL_NAMES, _SnapshotTarget, create_sampler_synthesis_tools, \
create_synthesis_tools, finalize_versioned_snapshot, \
make_write_snapshot_hook
SYNTHESIS_TOOL_NAMES, _SnapshotTarget, create_synthesis_tools, \
finalize_versioned_snapshot, make_write_snapshot_hook
from predicators.approaches.agent_bilevel_approach import AgentBilevelApproach
from predicators.approaches.sampler_learning_mixin import SamplerLearningMixin
from predicators.code_sim_learning.active_experiment import laplace_ensemble, \
mean_bernoulli_entropy, perturbation_ensemble, \
posterior_subsample_ensemble
Expand All @@ -45,12 +45,12 @@
iter_feature_residuals, merge_updates, read_latent_init, \
read_simulator_components
from predicators.envs import create_new_env
from predicators.ground_truth_models import get_gt_samplers, get_gt_simulator
from predicators.ground_truth_models import get_gt_simulator
from predicators.option_model import _OptionModelBase, _OracleOptionModel
from predicators.settings import CFG
from predicators.structs import Action, Dataset, GroundAtom, \
InteractionResult, LowLevelTrajectory, OptionSampler, \
ParameterizedOption, Predicate, State, Task, Type
InteractionResult, LowLevelTrajectory, ParameterizedOption, Predicate, \
State, Task, Type

logger = logging.getLogger(__name__)

Expand Down Expand Up @@ -78,7 +78,7 @@ def rule(state, updates, params):
# ── Approach ─────────────────────────────────────────────────────


class AgentSimLearningApproach(AgentBilevelApproach):
class AgentSimLearningApproach(SamplerLearningMixin, AgentBilevelApproach):
"""Bilevel planning with a learned step-level simulator.

During online learning:
Expand All @@ -92,6 +92,9 @@ class AgentSimLearningApproach(AgentBilevelApproach):
During solving:
- Uses the learned model for plan validation in backtracking
refinement.

Per-skill sampler learning (mode resolution, synthesis session
plumbing, loading) lives in :class:`SamplerLearningMixin`.
"""

def __init__(self,
Expand Down Expand Up @@ -152,13 +155,7 @@ def __init__(self,
# provenance (consumed in the next learn-phase prompt).
self._current_simulator_version: Optional[str] = None
self._current_predicates_version: Optional[str] = None
self._current_samplers_version: Optional[str] = None
# Whether this run learns samplers (vs. using ground-truth ones).
# Refined per cycle in _learn_simulator once GT availability is known;
# this default is what the synthesis-session tool surface reads.
self._do_synthesize_samplers: bool = (
CFG.agent_sim_learn_synthesize_samplers
and not CFG.agent_sim_learn_oracle_samplers)
self._init_sampler_learning_state()
# Partial-observability latent block: loaded from a simulator's
# LATENT_INIT export (None ⇒ no latent state). When the loaded
# rules use the recurrent 5-arg signature, fitting, the combined
Expand Down Expand Up @@ -305,296 +302,6 @@ def _build_synthesis_session_hooks(
],
}

# ── Per-skill sampler synthesis ─────────────────────────────
# Samplers are a first-class artifact of the base sim-learning
# approach (gated by a flag), not a subclass extension like
# predicates — so they are woven into _synthesize_with_agent and
# _learn_simulator directly rather than via the _extra_synthesis_*
# hooks, which keeps them independent of the predicate subclass's
# (non-super-calling) hook overrides. When a sim-synthesis session
# runs (oracle_sim_program=False) the sampler tool/snapshot/message
# ride along in it; when none runs (oracle_sim_program=True) they get
# a dedicated session via _synthesize_samplers_standalone.

@staticmethod
def _samplers_enabled() -> bool:
"""Whether per-skill samplers are used at all this run."""
return CFG.agent_sim_learn_synthesize_samplers

def _maybe_install_oracle_samplers(self) -> None:
"""Resolve sampler mode for this cycle and install GT ones if used.

Sets ``self._do_synthesize_samplers`` (learn vs. use ground
truth). When ``agent_sim_learn_oracle_samplers`` is on and the
env provides ground-truth samplers, installs them and skips
synthesis; if none exist, warns and falls back to synthesis.
"""
gt_samplers = None
if self._samplers_enabled() and CFG.agent_sim_learn_oracle_samplers:
gt_samplers = get_gt_samplers(CFG.env)
if gt_samplers:
self._synthesized_samplers = dict(gt_samplers)
self._current_samplers_version = "oracle"
logger.info("Using %d ground-truth sampler(s): %s",
len(gt_samplers), ", ".join(sorted(gt_samplers)))
else:
logger.warning(
"agent_sim_learn_oracle_samplers=True but no ground-truth "
"samplers for env %s; falling back to synthesis.", CFG.env)
self._do_synthesize_samplers = (self._samplers_enabled()
and not gt_samplers)

def _sampler_paths(self, base: str) -> Dict[str, str]:
"""Sandbox path bindings for samplers.py (host + agent-visible)."""
samplers_file = os.path.join(base, "samplers.py")
samplers_versions_dir = os.path.join(base, "samplers_versions")
if CFG.agent_sdk_use_local_sandbox:
samplers_file_for_agent = "./samplers.py"
elif self._tool_context.sandbox_dir:
samplers_file_for_agent = "/sandbox/samplers.py"
else:
samplers_file_for_agent = samplers_file
return {
"samplers_file": samplers_file,
"samplers_versions_dir": samplers_versions_dir,
"samplers_file_for_agent": samplers_file_for_agent,
}

def _make_sampler_tools(self, paths: Dict[str, str]) -> List[Any]:
"""Build the evaluate_sampler MCP tool for a synthesis session."""
return create_sampler_synthesis_tools(
samplers_file=paths["samplers_file"],
samplers_versions_dir=paths["samplers_versions_dir"],
approach=self,
cycle_index_provider=self._learning_cycle_index,
)

def _sampler_snapshot_target(self, paths: Dict[str,
str]) -> _SnapshotTarget:
"""Snapshot target that versions samplers.py on every Write/Edit."""
return _SnapshotTarget(
live_file=paths["samplers_file"],
versions_dir=paths["samplers_versions_dir"],
artifact_name="samplers",
cycle_index_provider=self._learning_cycle_index,
)

def _sampler_synthesis_message(self, paths: Dict[str, str]) -> str:
"""Instructions appended to the agent's first synthesis message."""
path = paths["samplers_file_for_agent"]
return f"""\
## Per-Skill Sampler Synthesis

Backtracking refinement draws each option's continuous parameters \
*uniformly* from its params box by default. When a sketch step's subgoal \
pins the parameters into a tiny region (e.g. a placement that must land \
within a few cm of an exact point and at a specific orientation), uniform \
sampling almost never hits it and refinement exhausts its budget. Fix this \
by writing per-skill samplers to `{path}` as a dict \
`LEARNED_SAMPLERS = {{"OptionName": sampler_fn, ...}}` keyed by option name.

Each sampler has signature \
`fn(state, subgoal_atoms, rng, objects) -> params` (the same signature as \
the env's NSRT samplers) where:
- `state` is the current `State` (read object features with `state.get(obj, "feat")`),
- `subgoal_atoms` is the set of `GroundAtom`s the step must establish — \
read the target relation here (e.g. an `InFront`/at-target atom names the \
two objects whose geometry the placement must satisfy) and compute the \
parameters that achieve it,
- `rng` is a `numpy` `Generator` (use it for small jitter so retries differ),
- `objects` is the list of typed objects bound to this option call.
Return a `float32` array whose length matches the option's params box \
(see `inspect_options` for the dimension and ranges); refinement clips it \
to that box, so stay within the ranges.

Aim the parameters at the subgoal geometrically (then add a little `rng` \
jitter); do NOT just return uniform draws. Read the option signatures with \
`inspect_options` and the predicate classifiers (for the subgoal geometry) \
with the predicate listing above.

Workflow: write `{path}`, call `evaluate_sampler` (snapshots + installs \
them and sanity-checks shape/box), then call `evaluate_plan_refinement` \
with a sketch using those options — the samples-to-refine count should \
drop sharply versus uniform. Iterate with `Edit` and re-run. Every \
successful Write/Edit of `{path}` is snapshotted to `samplers_versions/` \
as `cycle_XXX_vers_YYY_samplers.py`."""

def _finalize_and_load_samplers(self, paths: Dict[str, str]) -> None:
"""Snapshot the final samplers.py and load it into approach state."""
tag = finalize_versioned_snapshot(
paths["samplers_file"],
paths["samplers_versions_dir"],
cycle_idx=self._learning_cycle_index(),
artifact_name="samplers",
)
if tag is not None:
self._current_samplers_version = tag
logger.info("Final samplers snapshot: %s", tag)
loaded = self._load_samplers_from_module_file(paths["samplers_file"])
self._synthesized_samplers = loaded
logger.info("Loaded %d per-skill sampler(s) from %s.", len(loaded),
paths["samplers_file"])
for name in sorted(loaded):
logger.info(" sampler: %s", name)

def _load_samplers_from_module_file(self,
path: str) -> Dict[str, OptionSampler]:
"""Load LEARNED_SAMPLERS from ``path``; validate each entry.

Mirrors ``_load_predicates_from_module_file``. Returns an empty
dict on missing file or exec failure (samplers are optional).
Skips entries keyed by an unknown option name or whose value is
not callable.
"""
# pylint: disable=import-outside-toplevel
from predicators.agent_sdk.proposal_parser import build_exec_context, \
exec_code_safely
from predicators.agent_sdk.tools import _ParamsView

# pylint: enable=import-outside-toplevel
# ParamSpec is imported at module scope (used by exec'd samplers
# that close over learned params, mirroring the predicate loader).

if not os.path.isfile(path):
logger.info("No samplers file at %s; sampler set is empty.", path)
return {}

with open(path, "r", encoding="utf-8") as f:
code = f.read()

ctx = build_exec_context(types=self._types,
predicates=self._get_all_predicates(),
options=self._get_all_options(),
extra_context={
"params":
_ParamsView(self._fitted_params),
"ParamSpec": ParamSpec,
})

result, err = exec_code_safely(code, ctx, "LEARNED_SAMPLERS")
if err is not None:
logger.warning("Failed to load %s:\n%s", path, err)
return {}
if not isinstance(result, dict):
logger.warning("%s: LEARNED_SAMPLERS must be a dict, got %s.",
path,
type(result).__name__)
return {}

option_names = {o.name for o in self._get_all_options()}
valid: Dict[str, OptionSampler] = {}
for name, fn in result.items():
if name not in option_names:
logger.warning(
"Skipped sampler '%s' (not a known option name).", name)
continue
if not callable(fn):
logger.warning("Skipped sampler '%s' (value is not callable).",
name)
continue
valid[name] = fn
return valid

def _synthesize_samplers_standalone(
self, trajectories: List[LowLevelTrajectory],
base_pred_triples: List[Tuple[State, Action, State]],
inferred_hint: Dict[str, List[str]]) -> None:
"""Run a dedicated sampler-synthesis session.

Used when oracle_sim_program short-circuits the sim-synthesis
session, so samplers still get learned. Reuses that session's
sandbox/snapshot/tool machinery. Called from _learn_simulator
after the option model is built, so evaluate_plan_refinement has
a working simulator.
"""
if CFG.agent_sdk_use_local_sandbox:
sandbox_dir: Optional[str] = os.path.abspath(
os.path.join(self._get_log_dir(), "sandbox"))
else:
sandbox_dir = self._tool_context.sandbox_dir
base = sandbox_dir or self._get_log_dir()

if CFG.agent_sdk_use_local_sandbox:
sandbox_dir_for_agent: Optional[str] = "."
elif sandbox_dir:
sandbox_dir_for_agent = "/sandbox"
else:
sandbox_dir_for_agent = None

paths = self._sampler_paths(base)
simulator_file = os.path.join(base, "simulator.py")
versions_dir = os.path.join(base, "simulator_versions")

exec_ns: Dict[str, Any] = {
"trajectories":
trajectories,
"train_tasks":
self._train_tasks,
"is_goal_state":
lambda state, task_idx: self._train_tasks[task_idx].goal_holds(
state),
"np":
np,
"ParamSpec":
ParamSpec,
}
# evaluate_plan_refinement (from the standard synthesis tools) gives
# the agent the samples-to-refine feedback signal; the sampler tool
# installs + sanity-checks the samplers.
tools = create_synthesis_tools(
exec_ns,
base_pred_triples,
inferred_hint,
simulator_file=simulator_file,
versions_dir=versions_dir,
approach=self,
sandbox_dir=base,
sandbox_dir_for_agent=sandbox_dir_for_agent,
cycle_index_provider=self._learning_cycle_index,
)
tools.extend(self._make_sampler_tools(paths))
# Use the same declared surface as the mixin will assert against
# (_get_synthesis_tool_names already includes the sampler tool since
# _do_synthesize_samplers is True here). The rule-fitting tools are
# exposed but irrelevant — the message steers the agent to samplers.
declared = set(self._get_synthesis_tool_names() or ())
self._tool_context.extra_mcp_tools = [
t for t in tools if getattr(t, "name", "") in declared
]
self._learning_mode = True
self._tool_context.extra_session_hooks = (
self._build_synthesis_session_hooks(
[self._sampler_snapshot_target(paths)], base))

self._close_agent_session()
self._ensure_agent_session()

predicate_listing = self._format_predicate_signatures(
self._get_all_predicates())
message = f"""\
Synthesize per-skill samplers for this environment's options. The \
simulator dynamics are already fixed (oracle/learned); your only job is \
to make backtracking refinement land each option's continuous parameters \
on its sketch-step subgoal instead of drawing them uniformly.

## Available Predicates (subgoal geometry)
{predicate_listing}

Read the option signatures with `inspect_options` and explore the \
trajectory data with `run_python` (variables: `trajectories`, \
`train_tasks`, `is_goal_state`, `np`, `ParamSpec`)."""
message = message + "\n\n" + self._sampler_synthesis_message(paths)

try:
self._query_agent_sync(message, kind="learn")
finally:
self._tool_context.extra_session_hooks = {}
self._tool_context.extra_mcp_tools = []
self._learning_mode = False
self._close_agent_session()

self._finalize_and_load_samplers(paths)

# ── Learning ────────────────────────────────────────────────

def learn_from_offline_dataset(self, dataset: Dataset) -> None:
Expand Down
Loading
Loading