diff --git a/.github/workflows/docs.yml b/.github/workflows/docs.yml index 9e386814..56bf35ca 100644 --- a/.github/workflows/docs.yml +++ b/.github/workflows/docs.yml @@ -23,6 +23,8 @@ jobs: enable-cache: true - name: Install Python 3.13 run: uv python install 3.13 + - name: Generate API reference pages + run: uv run --python 3.13 --group dev python mujoco-mojo/scripts/gen_ref_pages.py - run: uv run --python 3.13 --group dev zensical build --clean - uses: actions/upload-pages-artifact@v4 with: diff --git a/.gitignore b/.gitignore index 9092d121..87398a38 100644 --- a/.gitignore +++ b/.gitignore @@ -13,6 +13,7 @@ mojo.log mojo.log.* site +docs/reference/ mojo-models textures diff --git a/docs/user-guides/getting-started.md b/docs/user-guides/getting-started.md index 59bf630c..1fb42c7a 100644 --- a/docs/user-guides/getting-started.md +++ b/docs/user-guides/getting-started.md @@ -306,5 +306,14 @@ MuJoCo does not currently ship first-party Python typing stubs. To enable proper --- +!!! tip "Tip: Scaffold a New Project" + Before diving into the generate script, run `mujoco-mojo init` in a new directory to get a working project skeleton — `simulation.py`, `run.sh`, and `reloaded.sh` — pre-wired and ready to edit. + + ```bash linenums="0" + mujoco-mojo init + ``` + + See the [Initializing a Project](init.md) guide for details. + !!! success This concludes our basic overview of the core concepts to MuJoCo Mojo. The next guide will cover how to get started with assembling your first model with the [generate script](generate-script.md). diff --git a/docs/user-guides/init.md b/docs/user-guides/init.md new file mode 100644 index 00000000..b5b36805 --- /dev/null +++ b/docs/user-guides/init.md @@ -0,0 +1,65 @@ +# Initializing a New Project + +!!! abstract + The `mujoco-mojo init` command scaffolds a new project in the current directory, providing ready-to-run starter files for Monte Carlo campaigns or optimization studies. + +--- + +## Usage + +```bash linenums="0" +mujoco-mojo init [--optimizer] +``` + +With no flags, the command generates a standard Monte Carlo project. Pass `--optimizer` to include an `objective` function stub for use with `mujoco-mojo run optimization`. + +--- + +## Generated Files + +Running `mujoco-mojo init` produces three files in the current directory: + +| File | Purpose | +|---|---| +| `simulation.py` | Simulation stubs: `UserData`, `generate`, and `runtime` | +| `run.sh` | Pre-configured shell script to launch the campaign | +| `reloaded.sh` | Shell script to open the interactive viewer | + +Both shell scripts are made executable and resolve their paths relative to their own location, so they can be invoked from any working directory. + +--- + +## Monte Carlo (default) + +```bash linenums="0" +mujoco-mojo init +``` + +`simulation.py` will contain stubs for `generate` and `runtime`. `run.sh` will call `mujoco-mojo run monte-carlo`. + +--- + +## Optimization + +```bash linenums="0" +mujoco-mojo init --optimizer +``` + +`simulation.py` will additionally contain an `objective` stub. `run.sh` will call `mujoco-mojo run optimization`. + +--- + +## Next Steps + +After running `init`, fill in your `generate` and `runtime` functions, then launch your campaign: + +```bash linenums="0" +bash run.sh +``` + +!!! tip + Use `reloaded.sh` during development to get rapid visual feedback on your model without running a full campaign. + + ```bash linenums="0" + bash reloaded.sh + ``` diff --git a/docs/user-guides/reloaded.md b/docs/user-guides/reloaded.md index 53250a5a..54e2da26 100644 --- a/docs/user-guides/reloaded.md +++ b/docs/user-guides/reloaded.md @@ -17,14 +17,20 @@ mujoco-mojo reloaded --generator my_sim.Experiment.generate The `reloaded` command is highly flexible, allowing you to specify different viewers, initial seeds, and even pass custom arguments to your scripts. -| Argument | Description | -|:-------------------|:--------------------------------------------------------------------------| -| `--help` | Describes all available arguments. Used on its own. | -| `--generator` | **(Required)** Import path to your `MojoGenerate` function. | -| `--runtime` | Optional import path to your `MojoRuntime` function. | -| `--user-interface` | Choose your viewer: `opengl` (native), `viser`, or `mjviser` (web-based). | -| `--seed` | Set the global seed for stochastic draws. | -| `--port` | Port for the web-based viewers (default: `8080`). | +| Argument | Shortcut | Description | +|:-------------------|:---------|:-------------------------------------------------------------------------------| +| `--help` | | Describes all available arguments. Used on its own. | +| `--generator` | `-g` | Import path to your `generate` function. | +| `--config` | `-c` | Path to a saved `model_config.json`. Mutually exclusive with `--generator`. | +| `--runtime` | `-r` | Optional import path to your `runtime` function. | +| `--workdir` | `-w` | Workspace directory for output files (default: `mojo-models`). | +| `--user-interface` | `-ui` | Choose your viewer: `opengl` (native), `viser`, or `mjviser` (web-based). | +| `--trial-num` | `-tn` | Load the random state for a specific trial number (default: `0`). | +| `--seed` | `-s` | Set the global seed for stochastic draws. | +| `--overrides` | `-o` | Path to a NamedValue overrides JSON file. Fixes specific distribution draws. | +| `--gen-arg` | `-ga` | Positional argument forwarded to the generator. Repeatable. | +| `--gen-kwarg` | `-gk` | Keyword argument (`key=value`) forwarded to the generator. Repeatable. | +| `--port` | `-p` | Port for the web-based viewers (default: `8080`). | --- @@ -41,6 +47,31 @@ Once Reloaded is running, your terminal becomes an interactive command center. Y --- +## Loading a Saved Model Config + +As an alternative to calling a generator function, Reloaded can load a previously saved `model_config.json` directly using `--config`. This is useful when you want to inspect the exact model that was used in a specific trial without running any Python generation logic. + +```bash linenums="0" +mujoco-mojo reloaded --config ./results/trial_0042/model_config.json +``` + +This mode is mutually exclusive with `--generator`. + +--- + +## Inspecting a Specific Trial + +When iterating on a stochastic simulation, it is often necessary to reproduce the exact random state of a particular trial. Pass `--trial-num` to seed the session with that trial's distribution draws. + +```bash linenums="0" +# Open trial 42 in the viewer, using the same seed and random draws as the original campaign +mujoco-mojo reloaded -g sim.generate -r sim.runtime --seed 123 --trial-num 42 +``` + +Combined with `--overrides`, you can also lock specific named values to fixed quantities while keeping everything else stochastic, which helps isolate the effect of individual parameters. + +--- + ## Workflow ### MJCF Prototyping diff --git a/docs/user-guides/running-jobs.md b/docs/user-guides/running-jobs.md index e905731e..7e0e6670 100644 --- a/docs/user-guides/running-jobs.md +++ b/docs/user-guides/running-jobs.md @@ -13,25 +13,23 @@ Mojo uses a single, unified command structure for all job types. By adjusting a ### Running a Single Trial -To run exactly one trial (usually trial `0`), simply set the trial count to `1` or specify a specific trial number. +Use `mujoco-mojo run single` when you want to execute exactly one trial. It accepts the same arguments as `run monte-carlo` and is the recommended command when running a baseline or re-running an individual trial for inspection. ```bash linenums="0" -# Run one trial using nominal values -mujoco-mojo run monte-carlo --generator sim.generate --runtime sim.runtime --seed 123 --n-trial 1 +# Run trial 0 (nominal values) +mujoco-mojo run single -g sim.generate -r sim.runtime --seed 123 -# Or run a specific trial from a previous distribution -mujoco-mojo run monte-carlo -g sim.generate -r sim.runtime --seed 123 --trial-num 42 +# Re-run a specific trial from a previous campaign +mujoco-mojo run single -g sim.generate -r sim.runtime --seed 123 --trial-num 42 ``` ???+ note "Note: Nominal Run" - For stochastic values, trial number `0` will use the distribution's `nominal_value` (if provided). This allows you to not only be able to recreate a model, but also build a "baseline" simulation. + Trial number `0` uses each distribution's `nominal_value` (if one was provided). This gives you a deterministic baseline that reflects design intent rather than a stochastic draw. - This helps you with defining the design intent, rather than letting chaos take over! - -More than one `--trial-num` argument can be passed to the command. This will make it so those select trials are run: +You can also target specific trials within a full `run monte-carlo` campaign by passing `--trial-num` one or more times. This is useful for re-running failed trials without re-running the entire job. ```bash linenums="0" -# This will run trials 12 and 42 +# Re-run only trials 12 and 42 mujoco-mojo run monte-carlo -g sim.generate -r sim.runtime --seed 123 --trial-num 12 --trial-num 42 ``` diff --git a/mkdocs.yml b/mkdocs.yml index 3fa84754..05075255 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -16,6 +16,7 @@ nav: - User Guide: # - user-guide.md - Getting Started: user-guides/getting-started.md + - Initializing a Project: user-guides/init.md - Generate Script: user-guides/generate-script.md - Runtime Script: user-guides/runtime-script.md - Running Jobs: @@ -27,7 +28,30 @@ nav: - Mosaic: user-guides/mosaic.md - Optimization: user-guides/optimization.md - Reloaded: user-guides/reloaded.md - # - Code Reference: reference/ + - Code Reference: + - reference/index.md + - mujoco_mojo: + - reference/mujoco_mojo/index.md + - MojoModel: reference/mujoco_mojo/mojo_model.md + - Base: reference/mujoco_mojo/base.md + - Visualization: reference/mujoco_mojo/visualization.md + - Typing: reference/mujoco_mojo/typing.md + # - Stochas: reference/mujoco_mojo/stochas.md + - MJCF: reference/mujoco_mojo/mjcf.md + - Runtime: + - reference/mujoco_mojo/runtime/index.md + - Load: reference/mujoco_mojo/runtime/load.md + - RuntimeManager: reference/mujoco_mojo/runtime/runtime_manager.md + - SignalManager: reference/mujoco_mojo/runtime/signal_manager.md + - VideoRecorder: reference/mujoco_mojo/runtime/video_recorder.md + - Utils: + - reference/mujoco_mojo/utils/index.md + - Color: reference/mujoco_mojo/utils/color.md + - Runner: reference/mujoco_mojo/utils/runner.md + - Proximity: reference/mujoco_mojo/utils/proximity.md + - Filters: reference/mujoco_mojo/utils/filters.md + - Interp: reference/mujoco_mojo/utils/interp.md + - Dataframe: reference/mujoco_mojo/utils/dataframe.md - About: about.md theme: @@ -81,12 +105,14 @@ plugins: - search - open-in-new-tab: add_icon: true + # gen-files and literate-nav are MkDocs-only plugins; zensical ignores them. + # Reference pages are pre-generated by scripts/gen_ref_pages.py instead. # - gen-files: # scripts: # - scripts/gen_ref_pages.py - - literate-nav: - nav_file: SUMMARY.md - - section-index + # - literate-nav: + # nav_file: SUMMARY.md + # - section-index - mkdocstrings: handlers: python: diff --git a/pyproject.toml b/pyproject.toml index f2fee582..0291f750 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,10 +1,10 @@ [build-system] -requires = ["setuptools>=61", "wheel"] -build-backend = "setuptools.build_meta" +requires = ["hatchling", "hatch-vcs"] +build-backend = "hatchling.build" [project] name = "mujoco-mojo" -version = "2.2.2" +dynamic = ["version"] description = "A complete MJCF lifecycle and trial orchestration suite for MuJoCo, powered by Pydantic v2." readme = "README.md" requires-python = ">=3.12" @@ -94,8 +94,12 @@ dev = [ [project.scripts] mujoco-mojo = "mujoco_mojo.utils.layers.cli:cli_app" -[tool.setuptools.package-data] -"*" = ["*.html", "*.js", "*.mp3", "*.svg", "*.css"] +[tool.hatch.version] +source = "vcs" +fallback-version = "0.0.0" + +[tool.hatch.build.targets.wheel] +packages = ["src/mujoco_mojo"] [tool.ruff] exclude = ["typings", "scripts/templating/sensor_template.py"] diff --git a/scripts/gen_ref_pages.py b/scripts/gen_ref_pages.py index 121f020a..34545849 100644 --- a/scripts/gen_ref_pages.py +++ b/scripts/gen_ref_pages.py @@ -1,35 +1,87 @@ -"""Generate the code reference pages and navigation.""" +""" +Generate API reference pages in docs/reference/. + +Run this script before building the docs site: + + python scripts/gen_ref_pages.py + +The output is written to docs/reference/ and is consumed by zensical (or +mkdocs) via the ::: autodoc directives. The directory is gitignored; this +script must be run as part of the CI pipeline before zensical build. +""" from pathlib import Path -import mkdocs_gen_files +ROOT = Path(__file__).parent.parent +DOCS_REF = ROOT / "docs" / "reference" + +# --------------------------------------------------------------------------- +# Modules to document: (dotted identifier, output path relative to DOCS_REF) +# --------------------------------------------------------------------------- +MODULES: list[tuple[str, str]] = [ + # top-level package + ("mujoco_mojo", "mujoco_mojo/index.md"), + ("mujoco_mojo.mojo_model", "mujoco_mojo/mojo_model.md"), + ("mujoco_mojo.base", "mujoco_mojo/base.md"), + ("mujoco_mojo.visualization", "mujoco_mojo/visualization.md"), + # typing — enums and type aliases (some not re-exported at top level) + ("mujoco_mojo.typing", "mujoco_mojo/typing.md"), + # stochas — distributions and named values (has items not in top-level __all__) + ("mujoco_mojo.stochas", "mujoco_mojo/stochas.md"), + # runtime + ("mujoco_mojo.runtime", "mujoco_mojo/runtime/index.md"), + ("mujoco_mojo.runtime.load", "mujoco_mojo/runtime/load.md"), + ("mujoco_mojo.runtime.runtime_manager", "mujoco_mojo/runtime/runtime_manager.md"), + ("mujoco_mojo.runtime.signal_manager", "mujoco_mojo/runtime/signal_manager.md"), + ("mujoco_mojo.runtime.video_recorder", "mujoco_mojo/runtime/video_recorder.md"), + # utils + ("mujoco_mojo.utils", "mujoco_mojo/utils/index.md"), + ("mujoco_mojo.utils.color", "mujoco_mojo/utils/color.md"), + ("mujoco_mojo.utils.runner", "mujoco_mojo/utils/runner.md"), + ("mujoco_mojo.utils.proximity", "mujoco_mojo/utils/proximity.md"), + ("mujoco_mojo.utils.filters", "mujoco_mojo/utils/filters.md"), + ("mujoco_mojo.utils.interp", "mujoco_mojo/utils/interp.md"), + ("mujoco_mojo.utils.dataframe", "mujoco_mojo/utils/dataframe.md"), +] + +# --------------------------------------------------------------------------- +# Hand-written stub for mujoco_mojo.mjcf (too large to autodoc) +# --------------------------------------------------------------------------- +MJCF_STUB = """\ +# mujoco\\_mojo.mjcf + +The `mujoco_mojo.mjcf` subpackage provides the complete Python object model +for the MuJoCo XML schema. It mirrors the official XML structure closely, so +the best reference is the upstream documentation: + +[MuJoCo XML Reference :material-open-in-new:](https://mujoco.readthedocs.io/en/stable/XMLreference.html){ .md-button .md-button--primary target=_blank } + +All `mjcf` types are re-exported from the top-level `mujoco_mojo` namespace, +so you rarely need to import directly from `mujoco_mojo.mjcf`. -nav = mkdocs_gen_files.Nav() # type: ignore +```python +import mujoco_mojo as mojo -root = Path(__file__).parent.parent -src = root / "src" +body = mojo.Body(name=mojo.BodyName("my_body")) +``` +""" -for path in sorted(src.rglob("*.py")): - module_path = path.relative_to(src).with_suffix("") - doc_path = path.relative_to(src).with_suffix(".md") - full_doc_path = Path("reference", doc_path) - parts = tuple(module_path.parts) +def main() -> None: + DOCS_REF.mkdir(parents=True, exist_ok=True) - if parts[-1] == "__init__": - parts = parts[:-1] - doc_path = doc_path.with_name("index.md") - full_doc_path = full_doc_path.with_name("index.md") - elif parts[-1] == "__main__": - continue + for module, rel_path in MODULES: + dest = DOCS_REF / rel_path + dest.parent.mkdir(parents=True, exist_ok=True) + dest.write_text(f"::: {module}\n", encoding="utf-8") + print(f" {dest.relative_to(ROOT)}") - nav[parts] = doc_path.as_posix() + mjcf_page = DOCS_REF / "mujoco_mojo" / "mjcf.md" + mjcf_page.write_text(MJCF_STUB, encoding="utf-8") + print(f" {mjcf_page.relative_to(ROOT)}") - with mkdocs_gen_files.open(full_doc_path, "w") as fd: - ident = ".".join(parts) - fd.write(f"::: {ident}") + print(f"\nWrote {len(MODULES) + 1} pages to {DOCS_REF.relative_to(ROOT)}/") - mkdocs_gen_files.set_edit_path(full_doc_path, path.relative_to(root)) -with mkdocs_gen_files.open("reference/SUMMARY.md", "w") as nav_file: - nav_file.writelines(nav.build_literate_nav()) +if __name__ == "__main__": + main() diff --git a/src/mujoco_mojo/templates/__init__.py b/src/mujoco_mojo/templates/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/src/mujoco_mojo/templates/monte_carlo.py b/src/mujoco_mojo/templates/monte_carlo.py new file mode 100644 index 00000000..a1963df6 --- /dev/null +++ b/src/mujoco_mojo/templates/monte_carlo.py @@ -0,0 +1,36 @@ +import mujoco + +import mujoco_mojo as mojo +import mujoco_mojo.runtime as rt + +logger = mojo.utils.get_logger(__name__) + + +class UserData(mojo.UserData): + """Pass state from generate() to runtime().""" + + pass + + +def generate(mojo_model: mojo.MojoModel, *args, **kwargs) -> mojo.MojoModel: + user_data = UserData() + mojo_model.mjcf.worldbody = mojo.WorldBody() + mojo_model.user_data = user_data + return mojo_model + + +def runtime( + mojo_model: mojo.MojoModel, + runtime_manager: rt.RuntimeManager, + mj_model: mujoco.MjModel, + mj_data: mujoco.MjData, + *args, + **kwargs, +) -> mojo.MojoModel: + _user_data = mojo_model.get_user_data(UserData) + + with runtime_manager as rm: + while mj_data.time < 1.0: + rm.step(mj_model, mj_data) + + return mojo_model diff --git a/src/mujoco_mojo/templates/optimization.py b/src/mujoco_mojo/templates/optimization.py new file mode 100644 index 00000000..5528b0d0 --- /dev/null +++ b/src/mujoco_mojo/templates/optimization.py @@ -0,0 +1,41 @@ +import mujoco + +import mujoco_mojo as mojo +import mujoco_mojo.runtime as rt + +logger = mojo.utils.get_logger(__name__) + + +class UserData(mojo.UserData): + """Pass state from generate() to runtime().""" + + pass + + +def generate(mojo_model: mojo.MojoModel, *args, **kwargs) -> mojo.MojoModel: + user_data = UserData() + mojo_model.mjcf.worldbody = mojo.WorldBody() + mojo_model.user_data = user_data + return mojo_model + + +def runtime( + mojo_model: mojo.MojoModel, + runtime_manager: rt.RuntimeManager, + mj_model: mujoco.MjModel, + mj_data: mujoco.MjData, + *args, + **kwargs, +) -> mojo.MojoModel: + _user_data = mojo_model.get_user_data(UserData) + + with runtime_manager as rm: + while mj_data.time < 1.0: + rm.step(mj_model, mj_data) + + return mojo_model + + +def objective(mojo_model: mojo.MojoModel, *args, **kwargs) -> float: + """Return a scalar score for the optimizer to minimize or maximize.""" + return 0.0 diff --git a/src/mujoco_mojo/templates/reloaded.sh b/src/mujoco_mojo/templates/reloaded.sh new file mode 100644 index 00000000..8417a852 --- /dev/null +++ b/src/mujoco_mojo/templates/reloaded.sh @@ -0,0 +1,7 @@ +#!/bin/bash +cd "$(dirname "${BASH_SOURCE[0]}")" + +mujoco-mojo reloaded \ + -g simulation.generate \ + -r simulation.runtime \ + -w ./results diff --git a/src/mujoco_mojo/templates/run_mc.sh b/src/mujoco_mojo/templates/run_mc.sh new file mode 100644 index 00000000..a702416a --- /dev/null +++ b/src/mujoco_mojo/templates/run_mc.sh @@ -0,0 +1,12 @@ +#!/bin/bash +cd "$(dirname "${BASH_SOURCE[0]}")" + +mujoco-mojo run monte-carlo \ + -g simulation.generate \ + -r simulation.runtime \ + -w ./results \ + -nt 100 \ + --no-resume \ + -cw \ + -np 4 \ + --seed 42 diff --git a/src/mujoco_mojo/templates/run_opt.sh b/src/mujoco_mojo/templates/run_opt.sh new file mode 100644 index 00000000..1b202a9f --- /dev/null +++ b/src/mujoco_mojo/templates/run_opt.sh @@ -0,0 +1,13 @@ +#!/bin/bash +cd "$(dirname "${BASH_SOURCE[0]}")" + +mujoco-mojo run optimization \ + -g simulation.generate \ + -r simulation.runtime \ + -ob simulation.objective \ + -w ./results \ + -nt 100 \ + --no-resume \ + -cw \ + -np 4 \ + --seed 42 diff --git a/src/mujoco_mojo/utils/layers/cli.py b/src/mujoco_mojo/utils/layers/cli.py index a0a59f67..48a049a0 100644 --- a/src/mujoco_mojo/utils/layers/cli.py +++ b/src/mujoco_mojo/utils/layers/cli.py @@ -721,7 +721,7 @@ def run_single( generator: GeneratorType, runtime: RuntimeType = DEFAULT_RUNTIME, workdir: WorkdirType = DEFAULT_WORKDIR, - n_trial: NTrialType = DEFAULT_MC_N_TRIAL, + n_trial: NTrialType = 1, n_proc: NProcType = DEFAULT_N_PROC, resume: ResumeType = DEFAULT_RESUME, seed: SeedType = DEFAULT_SEED, @@ -741,7 +741,7 @@ def run_single( """ [bold yellow]Execute a single trial.[/bold yellow] - This command handles the directory setup, distribution salting, and parallel execution of a single physics trial. + This command handles the directory setup, distribution salting, and execution of a single physics trial. """ from numpydantic import NDArray @@ -754,12 +754,12 @@ def run_single( workdir = workdir.resolve() - logger.info("Initializing Monte Carlo with CLI!") + logger.info("Initializing single trial with CLI!") dojo_cmd = f"mujoco-mojo dojo {workdir}" console.print( Panel( - "[bold green]Campaign Initialized![/]\n\n" + "[bold green]Trial Ready![/]\n\n" "[white]To monitor progress and view results, run:[/]\n" f" [bold yellow]{dojo_cmd}[/]", title="[cyan]Launch Control[/]", @@ -805,17 +805,12 @@ def run_single( run_kwargs=run_kwargs, ) - # 2. build config runner.config = MonteCarloConfig(n_trial=n_trial, n_proc=n_proc, resume=resume) - # 3. run - console.print( - f"[bold magenta]Starting {n_trial} trials[/bold magenta] (using {n_proc} workers)..." - ) - logger.info( - f"Starting {n_trial} trials (using {n_proc} workers)...", - extra={"file_only": True}, - ) + trial_id = trial_nums[0] if trial_nums else 0 + console.print(f"[bold magenta]Running trial {trial_id}[/bold magenta]...") + logger.info(f"Running trial {trial_id}...", extra={"file_only": True}) + had_fails = runner.run( global_overrides=global_overrides if global_overrides @@ -828,15 +823,15 @@ def run_single( match execution_mode: case ExecutionMode.LOCAL: if had_fails: - preamble = "[bold red]Monte Carlo finished with failures![/bold red]" + preamble = "[bold red]Trial finished with failures![/bold red]" logger.error( - f"Monte Carlo finished with failures! See results in {runner.workdir.resolve()}", + f"Trial finished with failures! See results in {runner.workdir.resolve()}", extra={"file_only": True}, ) else: - preamble = "[bold green]Monte Carlo finished![/bold green]" + preamble = "[bold green]Trial finished![/bold green]" logger.info( - f"Monte Carlo finished! See results in {runner.workdir.resolve()}", + f"Trial finished! See results in {runner.workdir.resolve()}", extra={"file_only": True}, ) console.print( @@ -844,19 +839,17 @@ def run_single( ) case ExecutionMode.SLURM: if had_fails: - finished_msg = ( - "[bold red]Failed to orchestrate SLURM Monte Carlo![/bold red]" - ) + finished_msg = "[bold red]Failed to orchestrate SLURM trial![/bold red]" logger.error( - "Failed to orchestrate SLURM Monte Carlo!", + "Failed to orchestrate SLURM trial!", extra={"file_only": True}, ) else: finished_msg = ( - "[bold green]SLURM Monte Carlo orchestration finished![/bold green]" + "[bold green]SLURM trial orchestration finished![/bold green]" ) logger.info( - "SLURM Monte Carlo orchestration finished!", + "SLURM trial orchestration finished!", extra={"file_only": True}, ) console.print(f"\n{finished_msg}") @@ -864,6 +857,72 @@ def run_single( raise typer.Exit() +@cli_app.command(name="init") +def init_project( + optimizer: Annotated[ + bool, + typer.Option( + "--optimizer/--no-optimizer", + "-op/-nop", + help="Include an [bold cyan]objective[/bold cyan] function stub for use with [bold]mujoco-mojo run optimization[/bold].", + ), + ] = False, +) -> None: + """ + [bold yellow]Initialize a new empty mujoco-mojo project.[/bold yellow] + + Writes a [bold cyan]simulation.py[/bold cyan] scaffold to the current directory containing stub implementations of [bold]generate[/bold], [bold]runtime[/bold], and [bold]UserData[/bold]. Pass [bold cyan]--optimizer[/bold cyan] to also include an [bold]objective[/bold] stub. + """ + import os + from importlib.resources import files + + tmpl = files("mujoco_mojo.templates") + + py_dest = Path("simulation.py") + run_dest = Path("run.sh") + reloaded_dest = Path("reloaded.sh") + + conflicts = [p for p in (py_dest, run_dest, reloaded_dest) if p.exists()] + if conflicts: + names = ", ".join(f"[bold cyan]{p}[/bold cyan]" for p in conflicts) + console.print( + f"[bold red]Error:[/bold red] {names} already exist. " + "Remove them first or rename them before running init." + ) + raise typer.Exit(code=1) + + py_template = "optimization.py" if optimizer else "monte_carlo.py" + sh_template = "run_opt.sh" if optimizer else "run_mc.sh" + + py_dest.write_text( + tmpl.joinpath(py_template).read_text(encoding="utf-8"), encoding="utf-8" + ) + run_dest.write_text( + tmpl.joinpath(sh_template).read_text(encoding="utf-8"), encoding="utf-8" + ) + reloaded_dest.write_text( + tmpl.joinpath("reloaded.sh").read_text(encoding="utf-8"), encoding="utf-8" + ) + + os.chmod(run_dest, 0o755) + os.chmod(reloaded_dest, 0o755) + + w = max(len(str(p)) for p in (py_dest, run_dest, reloaded_dest)) + console.print( + Panel( + f"[bold green]Project initialized![/bold green]\n\n" + f" [bold cyan]{str(py_dest).ljust(w)}[/bold cyan] - simulation stubs\n" + f" [bold cyan]{str(run_dest).ljust(w)}[/bold cyan] - run the campaign\n" + f" [bold cyan]{str(reloaded_dest).ljust(w)}[/bold cyan] - interactive viewer\n\n" + f"[white]Get started:[/white]\n" + f" [bold yellow]bash run.sh[/bold yellow]", + title="[cyan]Mojo Init[/cyan]", + expand=False, + border_style="cyan", + ) + ) + + @cli_app.command(name="reloaded") def run_reloaded( generator: ReloadedGeneratorType = None, diff --git a/src/mujoco_mojo/utils/layers/dojo/templates/base.html b/src/mujoco_mojo/utils/layers/dojo/templates/base.html index d677f527..562f4f2d 100644 --- a/src/mujoco_mojo/utils/layers/dojo/templates/base.html +++ b/src/mujoco_mojo/utils/layers/dojo/templates/base.html @@ -359,44 +359,40 @@

- -
+
-
- - + +
+
+ + +
+
- -
- - -
- - - - - + +
+ + +
+ + +
@@ -533,6 +529,50 @@

+
+ + + + + + +
+

+ {% block scripts %}{% endblock %} diff --git a/src/mujoco_mojo/utils/layers/dojo/templates/partials/trial_viewer/_header.html b/src/mujoco_mojo/utils/layers/dojo/templates/partials/trial_viewer/_header.html index f6a9ca66..98c00181 100644 --- a/src/mujoco_mojo/utils/layers/dojo/templates/partials/trial_viewer/_header.html +++ b/src/mujoco_mojo/utils/layers/dojo/templates/partials/trial_viewer/_header.html @@ -1,280 +1,280 @@ - -
- -
- - + +
+ + + + + All + + / + {{ trial_id }} +
+ +
+ +