def generate_construct(
self, use_runtime: bool, on_reload_callback: OnReloadCallback | None
) -> MjState:
"""Helper to run a generation."""
from .cli import _load_func
project_root = Path.cwd()
def reload_with_check(
path_str: str, label: str
) -> Callable[..., Any] | tuple[Callable[..., Any], types.ModuleType]:
"""Internal helper to validate file existance before reloading."""
try:
# attempt to find the module the dotted path resolves against, which may be a
# package `__init__` re-exporting a class/function defined in a submodule
_, mod = _load_func(path_str, _return_module=True)
if mod and hasattr(mod, "__file__") and mod.__file__:
source_file = Path(mod.__file__).resolve()
# edge case check: file no longer exists
if not source_file.exists():
self._print(
f"[bold red]Stop![/bold red] The source file for {label} is missing.\n"
f"Expected: [dim]{source_file}[/dim]\n"
f"[yellow]Hint:[/yellow] Did you rename [bold]'{source_file.name}'[/bold] or move it?"
)
raise FileNotFoundError(
f"Unable to find source file for {path_str}. Did you rename or move it?"
)
# reloading the entry module re-resolves any re-exported names (e.g.
# `from .robot import RobotGenerator`) after their source modules reload
recursive_reload(module=mod, project_root=project_root)
# reload to get the reference
return _load_func(path_str)
except (ModuleNotFoundError, ImportError) as e:
self._print(
f"[bold red]Error:[/bold red] Failed to import [bold green]{path_str}[/bold green]: {e}\n"
f"[yellow]Hint:[/yellow] Check the module path for typos, and check the imports inside your file."
)
raise
gen_func: MojoGenerator | None = None
if self.generator:
gen_func = cast(
MojoGenerator, reload_with_check(self.generator, "Generator")
)
run_func: MojoRuntime | None = None
if self.runtime:
run_func = cast(MojoRuntime, reload_with_check(self.runtime, "Runtime"))
global_overrides = NamedValueDict[NDArray]()
if self.overrides_path and self.overrides_path.exists():
logger.info(
f"Retrieving global NamedValue overrides from `{self.overrides_path}`"
)
global_overrides = NamedValueDict[NDArray].model_validate_json(
self.overrides_path.read_text()
)
if len(global_overrides) == 0:
logger.warning(
"Global NamedValue overrides had no entries. Continuing anyway."
)
else:
logger.info(
f"Global NamedValue overrides had {len(global_overrides)} entries."
)
# each trial number gets its own output folder, mirroring the Runner's layout
trial_dir = self.trial_dir_for(self.trial_num)
trial_dir.mkdir(parents=True, exist_ok=True)
self._current_trial_dir = trial_dir
# set up a per-trial log file capturing everything logged during this run, mirroring
# the Runner's mojo.log. Unlike the Runner, overwrite (rather than append) on each
# reload, since the same trial_dir is reused across many reload/run iterations in a
# single interactive session and appending would mix logs from unrelated runs.
trial_log_handler = get_trial_log_handler(trial_dir / "mojo.log", mode="w")
root_logger = logging.getLogger()
root_logger.addHandler(trial_log_handler)
# only worth tracking via TrialStatus/JobStatus if telemetry will actually be recorded
trial_status: TrialStatus | None = None
trial_status_path: Path | None = None
if self.record:
trial_status_path = trial_dir / TRIAL_STATUS_FNAME
trial_status = TrialStatus(trial_num=self.trial_num)
trial_status._path = trial_status_path
with trial_status.record_step(step_name="pending"):
pass
runtime_manager: rt.RuntimeManager | None = None
try:
# execute generation
start = time.time()
with self._record_step(trial_status, "generating"):
if self.generator:
assert gen_func
# explicitly reset all stochas registries, then reinsert overrides
mojo_model = MojoModel()
mojo_model.named = NamedValueDict[NDArray]()
mojo_model.dists = DistributionDict()
mojo_model.design = DesignValueDict()
mojo_model = (
mojo_model.with_overrides(overrides=global_overrides)
.with_seed(seed=self.seed)
.with_trial_num(self.trial_num)
)
mojo_model = gen_func(
mojo_model,
global_overrides,
*self.gen_args,
**self.gen_kwargs,
)
mojo_model._trial_dir = trial_dir
if self.model_config_name:
mojo_model.dump_to_path(trial_dir / self.model_config_name)
(trial_dir / NAMED_VALUES_FNAME).write_text(
mojo_model.named.model_dump_json()
)
# clear and rebuild on every reload (rather than write-once like
# the Runner) since the user's generator, and therefore which
# distributions exist (or whether any exist at all), can change
# between reloads; wiping first means a dropped category/dist_type
# or an emptied dists dict doesn't leave stale files behind
stochas_dir = self.workdir / STOCHAS_DIR_NAME
shutil.rmtree(stochas_dir, ignore_errors=True)
if mojo_model.dists:
stochas_dir.mkdir(exist_ok=True)
mojo_model.dists.to_tables(stochas_dir)
tmp = stochas_dir / "dists.tmp.json"
tmp.write_text(
mojo_model.dists.model_dump_json(), encoding="utf-8"
)
tmp.replace(stochas_dir / STOCHAS_DISTS_FNAME)
else:
assert self.config_path
try:
mojo_model = MojoModel.model_validate_json(
self.config_path.read_text()
).with_overrides(overrides=global_overrides)
mojo_model._trial_dir = trial_dir
logger.info(
f"Playback mode: Loaded frozen state from {self.config_path.name}"
)
except Exception as e:
msg = f"Failed to load model config: {e}"
logger.error(msg)
raise ValueError(msg)
try:
state = mojo_model.mjcf.prep_for_sim(
save_path=trial_dir / self.xml_name, unit_system=mojo_model.us
)
except Exception as e:
msg = f"Failed to compile with MuJoCo: {e}"
logger.error(msg)
raise ValueError(msg)
# sync the viewer
if on_reload_callback:
on_reload_callback(state)
self._print(
f"[dim white]Model generated in [bold]{time.time() - start:.2f}s[/bold].[/dim white]"
)
# execute runtime
with self._record_step(trial_status, "solving"):
if run_func and use_runtime:
runtime_manager = rt.RuntimeManager(
signal_manager=rt.SignalManager(
export_path=trial_dir
/ rt.SignalManager.default_output_name(),
unit_system=mojo_model.us,
),
_sync_hook=self._sync_hook,
_viewer_lock=self._viewer_lock,
_skip_recording=not self.record,
playback_speed=self._playback_speed,
_stop_event=self._stop_event,
)
start = time.time()
run_func(
mojo_model,
runtime_manager,
state,
*self.run_args,
**self.run_kwargs,
)
self._print(
f"[dim white]Runtime completed in [bold]{time.time() - start:.2f}s[/bold].[/dim white]"
)
else:
mujoco.mj_forward(state.model, state.data)
except rt.RequirementSatisfied:
# a live requirement ended the run early as a success; the outcome
# is decided by the requirement results, then the caller prints it
if trial_status is not None:
trial_status.step = "done"
if runtime_manager is not None:
trial_status.requirements = runtime_manager.requirement_results
trial_status.completion = (
Completion.FAILURE
if trial_status.requirements
and not all(r.passed for r in trial_status.requirements)
else Completion.SUCCESS
)
raise
except (BdbQuit, KeyboardInterrupt, rt.SimulationStopped):
raise
except Exception:
if trial_status is not None:
trial_status.step = "done"
trial_status.completion = Completion.ERROR
raise
else:
if trial_status is not None:
trial_status.step = "done"
trial_status.completion = Completion.SUCCESS
finally:
# persist the final trial status and fold it into the job-level tracker,
# the same way `MojoRunner` calls `update_trial` after each trial completes
if trial_status is not None and trial_status_path is not None:
trial_status.dump_to_path(trial_status_path)
job_status = self._ensure_job_status(gen_func, run_func)
if self.trial_num not in job_status.trial_nums:
job_status.trial_nums = sorted(
{*job_status.trial_nums, self.trial_num}
)
job_status.update_trial(status=trial_status)
root_logger.removeHandler(trial_log_handler)
trial_log_handler.close()
return state