From 12bffe7999c4f2c2674f28c25f5c456a1df6502d Mon Sep 17 00:00:00 2001 From: yuecideng Date: Sat, 23 May 2026 20:10:22 +0800 Subject: [PATCH] feat(toolkit): add embodichain-new-task CLI for task scaffolding Introduce a Jinja2-based scaffold that generates in-repo and extension task environments (demo, RL, config-only), with console entry point, package templates, and unit tests. Document the CLI in add-task-env skill. Co-authored-by: Cursor --- embodichain/toolkits/scaffold/__init__.py | 25 +++ embodichain/toolkits/scaffold/__main__.py | 22 ++ embodichain/toolkits/scaffold/cli.py | 191 ++++++++++++++++ embodichain/toolkits/scaffold/generator.py | 153 +++++++++++++ embodichain/toolkits/scaffold/naming.py | 58 +++++ embodichain/toolkits/scaffold/post_process.py | 127 +++++++++++ embodichain/toolkits/scaffold/presets.py | 208 +++++++++++++++++ embodichain/toolkits/scaffold/render.py | 77 +++++++ embodichain/toolkits/scaffold/spec.py | 212 ++++++++++++++++++ .../scaffold/templates/extension/README.md.j2 | 28 +++ .../templates/extension/constants.py.j2 | 21 ++ .../scaffold/templates/extension/gitignore.j2 | 9 + .../templates/extension/package_init.py.j2 | 35 +++ .../templates/extension/pyproject.toml.j2 | 34 +++ .../templates/extension/run_env.py.j2 | 46 ++++ .../extension/task_config_only.py.j2 | 28 +++ .../templates/extension/task_demo.py.j2 | 40 ++++ .../templates/extension/task_rl.py.j2 | 44 ++++ .../templates/extension/tasks_init.py.j2 | 19 ++ .../templates/extension/test_task.py.j2 | 24 ++ .../templates/inrepo/task_config_only.py.j2 | 33 +++ .../scaffold/templates/inrepo/task_demo.py.j2 | 45 ++++ .../scaffold/templates/inrepo/task_rl.py.j2 | 55 +++++ .../scaffold/templates/inrepo/test_task.py.j2 | 36 +++ pyproject.toml | 9 +- setup.py | 9 +- skills/add-task-env/SKILL.md | 24 ++ tests/toolkits/scaffold/test_scaffold.py | 115 ++++++++++ 28 files changed, 1725 insertions(+), 2 deletions(-) create mode 100644 embodichain/toolkits/scaffold/__init__.py create mode 100644 embodichain/toolkits/scaffold/__main__.py create mode 100644 embodichain/toolkits/scaffold/cli.py create mode 100644 embodichain/toolkits/scaffold/generator.py create mode 100644 embodichain/toolkits/scaffold/naming.py create mode 100644 embodichain/toolkits/scaffold/post_process.py create mode 100644 embodichain/toolkits/scaffold/presets.py create mode 100644 embodichain/toolkits/scaffold/render.py create mode 100644 embodichain/toolkits/scaffold/spec.py create mode 100644 embodichain/toolkits/scaffold/templates/extension/README.md.j2 create mode 100644 embodichain/toolkits/scaffold/templates/extension/constants.py.j2 create mode 100644 embodichain/toolkits/scaffold/templates/extension/gitignore.j2 create mode 100644 embodichain/toolkits/scaffold/templates/extension/package_init.py.j2 create mode 100644 embodichain/toolkits/scaffold/templates/extension/pyproject.toml.j2 create mode 100644 embodichain/toolkits/scaffold/templates/extension/run_env.py.j2 create mode 100644 embodichain/toolkits/scaffold/templates/extension/task_config_only.py.j2 create mode 100644 embodichain/toolkits/scaffold/templates/extension/task_demo.py.j2 create mode 100644 embodichain/toolkits/scaffold/templates/extension/task_rl.py.j2 create mode 100644 embodichain/toolkits/scaffold/templates/extension/tasks_init.py.j2 create mode 100644 embodichain/toolkits/scaffold/templates/extension/test_task.py.j2 create mode 100644 embodichain/toolkits/scaffold/templates/inrepo/task_config_only.py.j2 create mode 100644 embodichain/toolkits/scaffold/templates/inrepo/task_demo.py.j2 create mode 100644 embodichain/toolkits/scaffold/templates/inrepo/task_rl.py.j2 create mode 100644 embodichain/toolkits/scaffold/templates/inrepo/test_task.py.j2 create mode 100644 tests/toolkits/scaffold/test_scaffold.py diff --git a/embodichain/toolkits/scaffold/__init__.py b/embodichain/toolkits/scaffold/__init__.py new file mode 100644 index 00000000..5e35aa3e --- /dev/null +++ b/embodichain/toolkits/scaffold/__init__.py @@ -0,0 +1,25 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Scaffold new EmbodiChain task environments (in-repo or external extension).""" + +from __future__ import annotations + +from embodichain.toolkits.scaffold.cli import main +from embodichain.toolkits.scaffold.generator import generate_task +from embodichain.toolkits.scaffold.spec import TaskSpec + +__all__ = ["TaskSpec", "generate_task", "main"] diff --git a/embodichain/toolkits/scaffold/__main__.py b/embodichain/toolkits/scaffold/__main__.py new file mode 100644 index 00000000..bee8b2fe --- /dev/null +++ b/embodichain/toolkits/scaffold/__main__.py @@ -0,0 +1,22 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +from __future__ import annotations + +from embodichain.toolkits.scaffold.cli import main + +if __name__ == "__main__": + main() diff --git a/embodichain/toolkits/scaffold/cli.py b/embodichain/toolkits/scaffold/cli.py new file mode 100644 index 00000000..c0ce4ca5 --- /dev/null +++ b/embodichain/toolkits/scaffold/cli.py @@ -0,0 +1,191 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +from __future__ import annotations + +import argparse +import sys +from pathlib import Path + +from embodichain.toolkits.scaffold.generator import generate_task, print_summary +from embodichain.toolkits.scaffold.naming import default_gym_id +from embodichain.toolkits.scaffold.spec import INREPO_CATEGORIES, TaskSpec + + +def _build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser( + prog="embodichain-new-task", + description="Scaffold EmbodiChain task environments (in-repo or external extension).", + ) + parser.add_argument( + "--target", + choices=("inrepo", "extension"), + default="inrepo", + help="Generate inside EmbodiChain repo or a new extension project.", + ) + parser.add_argument( + "--workflow", + choices=("demo", "rl", "config-only"), + required=False, + help="Task workflow: expert demo, RL, or config-only env class.", + ) + parser.add_argument( + "--name", + "--task-name", + dest="task_snake", + help="Task name in snake_case (e.g. pick_place).", + ) + parser.add_argument("--gym-id", help="Gym registration id (e.g. PickPlace-v1).") + parser.add_argument( + "--category", + choices=INREPO_CATEGORIES, + default="tableware", + help="In-repo task category (ignored for RL workflow).", + ) + parser.add_argument( + "--robot-preset", + choices=("cobot_magic", "ur5_minimal"), + default="cobot_magic", + help="Robot/sensor/light preset for gym JSON.", + ) + parser.add_argument( + "--max-episode-steps", type=int, default=300, help="Max steps per episode." + ) + parser.add_argument( + "--max-episodes", type=int, default=5, help="Episodes in gym JSON metadata." + ) + parser.add_argument( + "--reward-style", + choices=("json", "python"), + default="json", + help="RL rewards in gym JSON or Python get_reward().", + ) + parser.add_argument( + "--project-name", + help="Extension project name (pyproject name).", + ) + parser.add_argument( + "--package-name", + help="Extension Python package name (snake_case).", + ) + parser.add_argument( + "--output-dir", + type=Path, + help="Extension output directory (default: ./).", + ) + parser.add_argument( + "--no-test", action="store_true", help="Skip generating test stub." + ) + parser.add_argument( + "--no-black", action="store_true", help="Skip running black on generated files." + ) + parser.add_argument( + "--init-git", + action="store_true", + help="Run git init in extension output (extension target only).", + ) + parser.add_argument( + "--force", action="store_true", help="Overwrite existing files." + ) + parser.add_argument( + "--dry-run", + action="store_true", + help="Print paths only; do not write files.", + ) + parser.add_argument( + "-i", + "--interactive", + action="store_true", + help="Prompt for missing options.", + ) + return parser + + +def _prompt_choice(label: str, options: list[str], default: str) -> str: + print(f"{label} [{'/'.join(options)}] (default: {default}): ", end="") + value = input().strip() + return value if value in options else default + + +def _interactive_fill(args: argparse.Namespace) -> None: + if args.target is None: + args.target = _prompt_choice("Target", ["inrepo", "extension"], "inrepo") + if args.workflow is None: + args.workflow = _prompt_choice( + "Workflow", ["demo", "rl", "config-only"], "demo" + ) + if args.task_snake is None: + args.task_snake = input("Task name (snake_case): ").strip() + if args.gym_id is None and args.task_snake: + args.gym_id = default_gym_id(args.task_snake) + print(f"Gym id (default: {args.gym_id}): ", end="") + custom = input().strip() + if custom: + args.gym_id = custom + if args.target == "inrepo" and args.workflow != "rl": + if args.category == "tableware": + args.category = _prompt_choice( + "Category", list(INREPO_CATEGORIES), "tableware" + ) + if args.target == "extension": + if args.package_name is None: + args.package_name = args.task_snake + if args.project_name is None: + args.project_name = args.package_name.replace("_", "-") + + +def main(argv: list[str] | None = None) -> int: + parser = _build_parser() + args = parser.parse_args(argv) + + if args.interactive or args.workflow is None or args.task_snake is None: + _interactive_fill(args) + + if args.workflow is None: + parser.error("--workflow is required (or use --interactive).") + if args.task_snake is None: + parser.error("--name is required (or use --interactive).") + + try: + spec = TaskSpec( + task_snake=args.task_snake, + workflow=args.workflow, + target=args.target, + gym_id=args.gym_id, + category=args.category, + robot_preset=args.robot_preset, + max_episode_steps=args.max_episode_steps, + max_episodes=args.max_episodes, + reward_style=args.reward_style, + project_name=args.project_name, + package_name=args.package_name, + output_dir=args.output_dir, + include_test=not args.no_test, + dry_run=args.dry_run, + force=args.force, + run_black=not args.no_black, + init_git=args.init_git, + ) + paths = generate_task(spec) + print_summary(spec, paths) + except (ValueError, FileExistsError) as exc: + print(f"Error: {exc}", file=sys.stderr) + return 1 + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/embodichain/toolkits/scaffold/generator.py b/embodichain/toolkits/scaffold/generator.py new file mode 100644 index 00000000..baf5f366 --- /dev/null +++ b/embodichain/toolkits/scaffold/generator.py @@ -0,0 +1,153 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +from __future__ import annotations + +import shutil +from pathlib import Path + +from embodichain.toolkits.scaffold import post_process +from embodichain.toolkits.scaffold.presets import gym_config_to_json +from embodichain.toolkits.scaffold.render import ( + render_extension_file, + render_task_py, + render_test_py, +) +from embodichain.toolkits.scaffold.spec import TaskSpec + +_APACHE_LICENSE = """Apache License +Version 2.0, January 2004 +http://www.apache.org/licenses/LICENSE-2.0 +""" + + +def _write(path: Path, content: str, spec: TaskSpec) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + if path.exists() and spec.force: + path.unlink() + path.write_text(content, encoding="utf-8") + + +def _collect_paths(spec: TaskSpec) -> list[Path]: + paths = list(spec.all_output_paths()) + if ( + spec.target == "extension" + and spec.output_dir is not None + and spec.is_new_extension_project() + and not spec.dry_run + ): + if ( + spec.output_dir.exists() + and any(spec.output_dir.iterdir()) + and not spec.force + ): + raise FileExistsError( + f"Output directory is not empty: {spec.output_dir} (use --force)" + ) + post_process.check_collisions(spec, paths) + return paths + + +def generate_task(spec: TaskSpec) -> list[Path]: + """Generate task scaffold files. Returns list of written paths.""" + _collect_paths(spec) + + if spec.dry_run: + return spec.all_output_paths() + + written: list[Path] = [] + + task_py = spec.task_py_path() + _write(task_py, render_task_py(spec), spec) + written.append(task_py) + + gym_json = spec.gym_json_path() + _write(gym_json, gym_config_to_json(spec), spec) + written.append(gym_json) + + if test_path := spec.test_py_path(): + _write(test_path, render_test_py(spec), spec) + written.append(test_path) + + if spec.target == "inrepo": + post_process.patch_tasks_init(spec) + elif spec.is_new_extension_project(): + _generate_extension_tree(spec, written) + else: + post_process.patch_tasks_init(spec) + + if spec.run_black: + post_process.run_black(written) + + if spec.target == "extension" and spec.init_git and spec.output_dir: + post_process.init_git_repo(spec.output_dir) + + return written + + +def _generate_extension_tree(spec: TaskSpec, written: list[Path]) -> None: + assert spec.output_dir is not None + assert spec.package_name is not None + root = spec.output_dir + pkg = root / spec.package_name + + license_text = _APACHE_LICENSE + repo_license = spec.repo_root / "LICENSE" + if repo_license.is_file(): + license_text = repo_license.read_text(encoding="utf-8") + + files: list[tuple[Path, str]] = [ + (root / "pyproject.toml", render_extension_file("pyproject.toml", spec)), + (root / "README.md", render_extension_file("README.md", spec)), + (root / "LICENSE", license_text), + (root / ".gitignore", render_extension_file("gitignore", spec)), + (root / "VERSION", "0.1.0\n"), + (pkg / "VERSION", "0.1.0\n"), + (pkg / "__init__.py", render_extension_file("package_init.py", spec)), + (pkg / "tasks" / "__init__.py", render_extension_file("tasks_init.py", spec)), + (pkg / "data" / "__init__.py", ""), + (pkg / "data" / "constants.py", render_extension_file("constants.py", spec)), + (pkg / "utils" / "__init__.py", ""), + (root / "scripts" / "run_env.py", render_extension_file("run_env.py", spec)), + ] + + for path, content in files: + _write(path, content, spec) + written.append(path) + + post_process.patch_tasks_init(spec) + + +def print_summary(spec: TaskSpec, paths: list[Path]) -> None: + """Print generation summary and next-step commands.""" + mode = "DRY RUN — would write" if spec.dry_run else "Wrote" + print(f"\n{mode} {len(paths)} file(s):\n") + for p in paths: + print(f" {p}") + + print("\nNext steps:\n") + if spec.target == "inrepo": + gym_cfg = spec.gym_json_path().relative_to(spec.repo_root) + print(f" python embodichain/lab/scripts/run_env.py " f"--gym_config {gym_cfg}") + print(" pytest " + str(spec.test_py_path().relative_to(spec.repo_root))) + else: + assert spec.output_dir is not None + print(f" cd {spec.output_dir}") + print(" pip install -e .") + print( + f" python scripts/run_env.py " + f"--gym_config configs/{spec.task_snake}/gym.json --headless" + ) diff --git a/embodichain/toolkits/scaffold/naming.py b/embodichain/toolkits/scaffold/naming.py new file mode 100644 index 00000000..6d2de176 --- /dev/null +++ b/embodichain/toolkits/scaffold/naming.py @@ -0,0 +1,58 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +from __future__ import annotations + +import re + +_SNAKE_RE = re.compile(r"^[a-z][a-z0-9_]*$") + + +def snake_to_camel(snake: str) -> str: + """Convert ``pick_place`` to ``PickPlace``.""" + return "".join(part.capitalize() for part in snake.split("_") if part) + + +def default_env_class(snake: str) -> str: + """Convert ``pick_place`` to ``PickPlaceEnv``.""" + return f"{snake_to_camel(snake)}Env" + + +def default_gym_id(snake: str, *, version: str = "v1") -> str: + """Default gym id from task snake name.""" + return f"{snake_to_camel(snake)}-{version}" + + +def validate_snake(name: str) -> str: + if not _SNAKE_RE.match(name): + raise ValueError( + f"Task name must be snake_case (e.g. pick_place), got: {name!r}" + ) + return name + + +def validate_gym_id(gym_id: str) -> str: + if not gym_id or " " in gym_id: + raise ValueError(f"Invalid gym id: {gym_id!r}") + return gym_id + + +def validate_package_name(name: str) -> str: + if not re.match(r"^[a-z][a-z0-9_]*$", name): + raise ValueError( + f"Package name must be a valid Python module name, got: {name!r}" + ) + return name diff --git a/embodichain/toolkits/scaffold/post_process.py b/embodichain/toolkits/scaffold/post_process.py new file mode 100644 index 00000000..039f8641 --- /dev/null +++ b/embodichain/toolkits/scaffold/post_process.py @@ -0,0 +1,127 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +from __future__ import annotations + +import ast +import subprocess +import sys +from pathlib import Path + +from embodichain.toolkits.scaffold.spec import TaskSpec + +_MARKER_PREFIX = "# --- embodichain-new-task:" + + +def check_collisions(spec: TaskSpec, paths: list[Path]) -> None: + """Raise if any output path exists and ``force`` is False.""" + existing = [p for p in paths if p.exists()] + if existing and not spec.force: + lines = "\n ".join(str(p) for p in existing) + raise FileExistsError( + f"Output path(s) already exist (use --force to overwrite):\n {lines}" + ) + + +def patch_tasks_init(spec: TaskSpec) -> None: + """Append import and ``__all__`` entry to tasks ``__init__.py``.""" + init_path = spec.tasks_init_path() + if not init_path.exists(): + init_path.parent.mkdir(parents=True, exist_ok=True) + init_path.write_text( + "# Auto-generated tasks package.\nfrom __future__ import annotations\n\n" + "__all__ = []\n", + encoding="utf-8", + ) + + content = init_path.read_text(encoding="utf-8") + marker = f"{_MARKER_PREFIX} {spec.gym_id}" + import_line = spec.tasks_init_import_line() + + if import_line in content: + content = _upsert_all(content, spec.task_class) + init_path.write_text(content, encoding="utf-8") + return + + if import_line not in content: + if not content.endswith("\n"): + content += "\n" + content += f"\n{marker}\n{import_line}\n" + + content = _upsert_all(content, spec.task_class) + init_path.write_text(content, encoding="utf-8") + + +def _upsert_all(content: str, class_name: str) -> str: + try: + tree = ast.parse(content) + except SyntaxError: + if f'"{class_name}"' not in content: + if "__all__" in content: + content = content.replace( + "__all__ = [", + f'__all__ = [\n "{class_name}",', + 1, + ) + else: + content += f'\n__all__ = ["{class_name}"]\n' + return content + + for node in tree.body: + if isinstance(node, ast.Assign): + for target in node.targets: + if isinstance(target, ast.Name) and target.id == "__all__": + if isinstance(node.value, ast.List): + elts = node.value.elts + names = [] + for elt in elts: + if isinstance(elt, ast.Constant) and isinstance( + elt.value, str + ): + names.append(elt.value) + if class_name not in names: + names.append(class_name) + new_list = ", ".join(repr(n) for n in names) + lines = content.splitlines() + start = node.lineno - 1 + end = node.end_lineno or node.lineno + lines[start:end] = [f"__all__ = [{new_list}]"] + return "\n".join(lines) + ( + "\n" if content.endswith("\n") else "" + ) + if "__all__" not in content: + content += f'\n__all__ = ["{class_name}"]\n' + return content + + +def run_black(paths: list[Path]) -> None: + py_files = [str(p) for p in paths if p.suffix == ".py" and p.exists()] + if not py_files: + return + subprocess.run( + [sys.executable, "-m", "black", *py_files], + check=False, + ) + + +def init_git_repo(root: Path) -> None: + subprocess.run(["git", "init"], cwd=root, check=False) + subprocess.run(["git", "add", "-f", "."], cwd=root, check=False) + subprocess.run( + ["git", "commit", "-q", "-m", "Initial commit from embodichain-new-task"], + cwd=root, + check=False, + ) diff --git a/embodichain/toolkits/scaffold/presets.py b/embodichain/toolkits/scaffold/presets.py new file mode 100644 index 00000000..8b9c0e1a --- /dev/null +++ b/embodichain/toolkits/scaffold/presets.py @@ -0,0 +1,208 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +from __future__ import annotations + +import json +from typing import Any + +from embodichain.toolkits.scaffold.spec import RobotPreset, TaskSpec + +ROBOT_PRESETS: dict[RobotPreset, dict[str, Any]] = { + "cobot_magic": { + "robot": { + "uid": "CobotMagic", + "robot_type": "CobotMagic", + "init_pos": [0.0, 0.0, 0.7775], + "init_qpos": [ + -0.3, + 0.3, + 1.0, + 1.0, + -1.2, + -1.2, + 0.0, + 0.0, + 0.6, + 0.6, + 0.0, + 0.0, + 0.05, + 0.05, + 0.05, + 0.05, + ], + }, + "sensor": [ + { + "sensor_type": "Camera", + "uid": "cam_high", + "width": 640, + "height": 480, + "intrinsics": [ + 488.1665344238281, + 488.1665344238281, + 322.7323303222656, + 213.17434692382812, + ], + "extrinsics": {"eye": [1.0, 0.0, 2.0], "target": [0.0, 0.0, 1.0]}, + } + ], + "light": { + "indirect": { + "emission_light": {"color": [1.0, 1.0, 1.0], "intensity": 150} + }, + "direct": [ + { + "uid": "light_1", + "light_type": "point", + "color": [1.0, 1.0, 1.0], + "intensity": 20.0, + "init_pos": [0.0, 0.0, 3.0], + "radius": 10.0, + } + ], + }, + }, + "ur5_minimal": { + "robot": { + "uid": "UR5", + "fpath": "UniversalRobots/UR5/UR5.urdf", + "init_pos": [0.0, 0.0, 0.7775], + "init_qpos": [ + 1.57079, + -1.57079, + 1.57079, + -1.57079, + -1.57079, + -3.14159, + ], + }, + "sensor": [ + { + "sensor_type": "Camera", + "uid": "cam_high", + "width": 640, + "height": 480, + "intrinsics": [488.1665344238281, 488.1665344238281, 320.0, 240.0], + "extrinsics": {"eye": [1.0, 0.0, 3.0], "target": [0.0, 0.0, 1.0]}, + } + ], + "light": { + "indirect": { + "emission_light": {"color": [1.0, 1.0, 1.0], "intensity": 150} + }, + "direct": [ + { + "uid": "light_1", + "light_type": "point", + "color": [1.0, 1.0, 1.0], + "intensity": 20.0, + "init_pos": [0.0, 0.0, 3.0], + "radius": 10.0, + } + ], + }, + }, +} + + +def build_gym_config(spec: TaskSpec) -> dict[str, Any]: + """Build a minimal gym JSON config for the task.""" + preset = ROBOT_PRESETS[spec.robot_preset] + env_block: dict[str, Any] = {"events": {}, "observations": {}} + + if spec.workflow == "demo": + env_block["events"]["record_camera"] = { + "func": "record_camera_data", + "mode": "interval", + "interval_step": 1, + "params": { + "name": "cam1", + "resolution": [320, 240], + "eye": [2.0, 0.0, 2.0], + "target": [0.5, 0.0, 1.0], + }, + } + elif spec.workflow == "rl" and spec.reward_style == "json": + env_block["events"] = { + "randomize_object": { + "func": "randomize_rigid_object_pose", + "mode": "reset", + "params": { + "entity_cfg": {"uid": "object"}, + "position_range": [[-0.1, -0.1, 0.0], [0.1, 0.1, 0.0]], + "relative_position": True, + }, + } + } + env_block["observations"] = { + "robot_qpos": { + "func": "normalize_robot_joint_data", + "mode": "modify", + "name": "robot/qpos", + "params": {}, + } + } + env_block["rewards"] = { + "task_reward": { + "func": "success_reward", + "mode": "add", + "weight": 1.0, + "params": {}, + } + } + env_block["actions"] = { + "delta_qpos": { + "func": "DeltaQposTerm", + "params": {"scale": 0.1}, + } + } + env_block["extensions"] = {"success_threshold": 0.05} + + config: dict[str, Any] = { + "id": spec.gym_id, + "max_episodes": spec.max_episodes, + "max_episode_steps": spec.max_episode_steps, + "env": env_block, + "robot": preset["robot"], + "sensor": preset["sensor"], + "light": preset["light"], + "background": [], + "rigid_object": [], + "rigid_object_group": [], + "articulation": [], + } + + if spec.workflow == "rl": + config["rigid_object"] = [ + { + "uid": "object", + "shape": {"type": "box", "size": [0.05, 0.05, 0.05]}, + "physical_attr": { + "mass": 0.1, + "static_friction": 0.5, + "dynamic_friction": 0.5, + }, + "init_pos": [0.5, 0.0, 0.05], + } + ] + + return config + + +def gym_config_to_json(spec: TaskSpec, *, indent: int = 4) -> str: + return json.dumps(build_gym_config(spec), indent=indent) + "\n" diff --git a/embodichain/toolkits/scaffold/render.py b/embodichain/toolkits/scaffold/render.py new file mode 100644 index 00000000..c11e51b2 --- /dev/null +++ b/embodichain/toolkits/scaffold/render.py @@ -0,0 +1,77 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +from __future__ import annotations + +from pathlib import Path + +from jinja2 import Environment, FileSystemLoader, select_autoescape + +from embodichain.toolkits.scaffold.presets import gym_config_to_json +from embodichain.toolkits.scaffold.spec import TaskSpec + +_TEMPLATE_DIR = Path(__file__).resolve().parent / "templates" + +_env = Environment( + loader=FileSystemLoader(_TEMPLATE_DIR), + autoescape=select_autoescape(enabled_extensions=()), + trim_blocks=True, + lstrip_blocks=True, + keep_trailing_newline=True, +) + + +def _ctx(spec: TaskSpec) -> dict: + return { + "spec": spec, + "task_snake": spec.task_snake, + "task_class": spec.task_class, + "gym_id": spec.gym_id, + "description": spec.description, + "max_episode_steps": spec.max_episode_steps, + "package_name": spec.package_name or "", + "project_name": spec.project_name or "", + "gym_json": gym_config_to_json(spec), + } + + +def render_task_py(spec: TaskSpec) -> str: + if spec.workflow == "rl": + template = "inrepo/task_rl.py.j2" + if spec.target == "extension": + template = "extension/task_rl.py.j2" + elif spec.workflow == "config-only": + template = "inrepo/task_config_only.py.j2" + if spec.target == "extension": + template = "extension/task_config_only.py.j2" + else: + template = "inrepo/task_demo.py.j2" + if spec.target == "extension": + template = "extension/task_demo.py.j2" + return _env.get_template(template).render(**_ctx(spec)) + + +def render_test_py(spec: TaskSpec) -> str: + template = ( + "extension/test_task.py.j2" + if spec.target == "extension" + else "inrepo/test_task.py.j2" + ) + return _env.get_template(template).render(**_ctx(spec)) + + +def render_extension_file(name: str, spec: TaskSpec) -> str: + return _env.get_template(f"extension/{name}.j2").render(**_ctx(spec)) diff --git a/embodichain/toolkits/scaffold/spec.py b/embodichain/toolkits/scaffold/spec.py new file mode 100644 index 00000000..e05cd44d --- /dev/null +++ b/embodichain/toolkits/scaffold/spec.py @@ -0,0 +1,212 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +from __future__ import annotations + +from dataclasses import dataclass, field +from pathlib import Path +from typing import Literal + +from embodichain.toolkits.scaffold.naming import ( + default_env_class, + default_gym_id, + validate_gym_id, + validate_package_name, + validate_snake, +) + +Target = Literal["inrepo", "extension"] +Workflow = Literal["demo", "rl", "config-only"] +InrepoCategory = Literal["tableware", "rl", "special"] +RobotPreset = Literal["cobot_magic", "ur5_minimal"] + + +INREPO_CATEGORIES = ("tableware", "rl", "special") + + +@dataclass +class TaskSpec: + """Specification for generating a new task environment.""" + + task_snake: str + workflow: Workflow + target: Target = "inrepo" + + gym_id: str | None = None + task_class: str | None = None + description: str = "" + + category: InrepoCategory = "tableware" + project_name: str | None = None + package_name: str | None = None + output_dir: Path | None = None + + robot_preset: RobotPreset = "cobot_magic" + max_episode_steps: int = 300 + max_episodes: int = 5 + num_envs: int = 1 + + with_rl_config: bool = False + reward_style: Literal["json", "python"] = "json" + include_test: bool = True + dry_run: bool = False + force: bool = False + run_black: bool = True + init_git: bool = False + + repo_root: Path = field(default_factory=lambda: Path(__file__).resolve().parents[3]) + + def __post_init__(self) -> None: + self.task_snake = validate_snake(self.task_snake) + if self.target == "inrepo" and self.category not in INREPO_CATEGORIES: + raise ValueError( + f"category must be one of {INREPO_CATEGORIES}, got {self.category!r}" + ) + if self.workflow == "rl" and self.target == "inrepo": + self.category = "rl" + if self.gym_id is None: + self.gym_id = default_gym_id(self.task_snake) + else: + self.gym_id = validate_gym_id(self.gym_id) + if self.task_class is None: + self.task_class = default_env_class(self.task_snake) + else: + if not self.task_class.endswith("Env"): + raise ValueError( + f"task_class should end with 'Env', got {self.task_class!r}" + ) + if not self.description: + self.description = f"{self.task_class} task environment (generated by embodichain-new-task)." + if self.target == "extension": + if self.package_name is None: + self.package_name = self.task_snake + validate_package_name(self.package_name) + if self.project_name is None: + self.project_name = self.package_name.replace("_", "-") + if self.output_dir is None: + self.output_dir = Path.cwd() / self.project_name + else: + self.output_dir = Path(self.output_dir).resolve() + else: + self.output_dir = self.repo_root + + if self.workflow == "rl": + self.with_rl_config = True + + @property + def module_name(self) -> str: + return self.task_snake + + def task_py_path(self) -> Path: + if self.target == "extension": + assert self.package_name is not None + return ( + self.output_dir / self.package_name / "tasks" / f"{self.task_snake}.py" + ) + return ( + self.repo_root + / "embodichain" + / "lab" + / "gym" + / "envs" + / "tasks" + / self.category + / f"{self.task_snake}.py" + ) + + def gym_json_path(self) -> Path: + if self.target == "extension": + return self.output_dir / "configs" / self.task_snake / "gym.json" + if self.workflow == "rl": + return ( + self.repo_root + / "configs" + / "agents" + / "rl" + / self.task_snake + / "gym_config.json" + ) + return self.repo_root / "configs" / "gym" / self.task_snake / "gym.json" + + def test_py_path(self) -> Path | None: + if not self.include_test: + return None + if self.target == "extension": + return self.output_dir / "tests" / f"test_{self.task_snake}.py" + return ( + self.repo_root + / "tests" + / "gym" + / "envs" + / "tasks" + / f"test_{self.task_snake}.py" + ) + + def tasks_init_path(self) -> Path: + if self.target == "extension": + assert self.package_name is not None + return self.output_dir / self.package_name / "tasks" / "__init__.py" + return ( + self.repo_root + / "embodichain" + / "lab" + / "gym" + / "envs" + / "tasks" + / "__init__.py" + ) + + def is_new_extension_project(self) -> bool: + if self.target != "extension" or self.output_dir is None: + return False + return not (self.output_dir / "pyproject.toml").exists() + + def all_output_paths(self) -> list[Path]: + paths = [self.task_py_path(), self.gym_json_path()] + if test := self.test_py_path(): + paths.append(test) + if self.target == "extension" and self.is_new_extension_project(): + assert self.package_name is not None + pkg = self.output_dir / self.package_name + paths.extend( + [ + pkg / "__init__.py", + pkg / "VERSION", + self.output_dir / "pyproject.toml", + self.output_dir / "README.md", + self.output_dir / "LICENSE", + self.output_dir / ".gitignore", + self.output_dir / "VERSION", + self.output_dir / "scripts" / "run_env.py", + pkg / "data" / "__init__.py", + pkg / "data" / "constants.py", + pkg / "utils" / "__init__.py", + pkg / "tasks" / "__init__.py", + ] + ) + return paths + + def import_path(self) -> str: + if self.target == "extension": + assert self.package_name is not None + return f"{self.package_name}.tasks.{self.task_snake}" + return f"embodichain.lab.gym.envs.tasks.{self.category}.{self.task_snake}" + + def tasks_init_import_line(self) -> str: + """Import line written into ``tasks/__init__.py``.""" + if self.target == "extension": + return f"from .{self.task_snake} import {self.task_class}" + return f"from {self.import_path()} import {self.task_class}" diff --git a/embodichain/toolkits/scaffold/templates/extension/README.md.j2 b/embodichain/toolkits/scaffold/templates/extension/README.md.j2 new file mode 100644 index 00000000..412f0647 --- /dev/null +++ b/embodichain/toolkits/scaffold/templates/extension/README.md.j2 @@ -0,0 +1,28 @@ +# {{ project_name }} + +EmbodiChain external task package generated by `embodichain-new-task`. + +## Installation + +Install [EmbodiChain](https://github.com/DexForce/EmbodiChain) first, then: + +```bash +pip install -e . +``` + +## Run task + +```bash +python scripts/run_env.py --gym_config configs/{{ task_snake }}/gym.json --headless +``` + +Preview mode: + +```bash +python scripts/run_env.py --gym_config configs/{{ task_snake }}/gym.json --preview +``` + +## Project layout + +- `{{ package_name }}/tasks/` — task env classes (`@register_env`) +- `configs/{{ task_snake }}/gym.json` — scene and manager configuration diff --git a/embodichain/toolkits/scaffold/templates/extension/constants.py.j2 b/embodichain/toolkits/scaffold/templates/extension/constants.py.j2 new file mode 100644 index 00000000..f7090144 --- /dev/null +++ b/embodichain/toolkits/scaffold/templates/extension/constants.py.j2 @@ -0,0 +1,21 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +from __future__ import annotations + +import os + +PROJECT_DIR = os.path.dirname(os.path.dirname(__file__)) diff --git a/embodichain/toolkits/scaffold/templates/extension/gitignore.j2 b/embodichain/toolkits/scaffold/templates/extension/gitignore.j2 new file mode 100644 index 00000000..66fa4073 --- /dev/null +++ b/embodichain/toolkits/scaffold/templates/extension/gitignore.j2 @@ -0,0 +1,9 @@ +__pycache__/ +*.py[cod] +*.egg-info/ +.eggs/ +dist/ +build/ +.pytest_cache/ +.mypy_cache/ +*.so diff --git a/embodichain/toolkits/scaffold/templates/extension/package_init.py.j2 b/embodichain/toolkits/scaffold/templates/extension/package_init.py.j2 new file mode 100644 index 00000000..a6e6921a --- /dev/null +++ b/embodichain/toolkits/scaffold/templates/extension/package_init.py.j2 @@ -0,0 +1,35 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +from __future__ import annotations + +import os + +project_dir = os.path.dirname(__file__) + + +def _get_version() -> str: + version_file = os.path.join(project_dir, "VERSION") + try: + with open(version_file, encoding="utf-8") as f: + return f.read().strip() + except FileNotFoundError: + return "0.1.0" + + +__version__ = _get_version() + +import {{ package_name }}.tasks # noqa: E402, F401 diff --git a/embodichain/toolkits/scaffold/templates/extension/pyproject.toml.j2 b/embodichain/toolkits/scaffold/templates/extension/pyproject.toml.j2 new file mode 100644 index 00000000..c3299034 --- /dev/null +++ b/embodichain/toolkits/scaffold/templates/extension/pyproject.toml.j2 @@ -0,0 +1,34 @@ +[project] +name = "{{ project_name }}" +description = "EmbodiChain task extension: {{ project_name }}" +readme = "README.md" +requires-python = ">=3.10" +dynamic = ["version"] +license = {text = "Apache-2.0"} +authors = [ + {name = "To be determined"} +] +keywords = ["robotics", "embodied-ai", "manipulation", "simulation"] +classifiers = [ + "Development Status :: 3 - Alpha", + "License :: OSI Approved :: Apache Software License", + "Programming Language :: Python :: 3", +] + +dependencies = [] + +[project.urls] +Repository = "https://github.com/DexForce/embodichain_task_template" + +[build-system] +requires = ["setuptools>=61.0", "wheel"] +build-backend = "setuptools.build_meta" + +[tool.setuptools.packages.find] +where = ["."] +exclude = ["docs", "tests"] + +[tool.setuptools.dynamic] +version = { file = ["{{ package_name }}/VERSION"] } + +[tool.black] diff --git a/embodichain/toolkits/scaffold/templates/extension/run_env.py.j2 b/embodichain/toolkits/scaffold/templates/extension/run_env.py.j2 new file mode 100644 index 00000000..8b64ceb0 --- /dev/null +++ b/embodichain/toolkits/scaffold/templates/extension/run_env.py.j2 @@ -0,0 +1,46 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Script to run the environment.""" + +from __future__ import annotations + +import argparse + +import gymnasium as gym +import numpy as np +import torch + +import {{ package_name }} # noqa: F401 — registers task environments + +from embodichain.lab.gym.utils.gym_utils import ( + add_env_launcher_args_to_parser, + build_env_cfg_from_args, +) +from embodichain.lab.scripts.run_env import main as run_env_main + + +if __name__ == "__main__": + np.set_printoptions(precision=5, suppress=True) + torch.set_printoptions(precision=5, sci_mode=False) + + parser = argparse.ArgumentParser() + add_env_launcher_args_to_parser(parser) + args = parser.parse_args() + + env_cfg, gym_config, action_config = build_env_cfg_from_args(args) + env = gym.make(id=gym_config["id"], cfg=env_cfg, **action_config) + run_env_main(args, env, gym_config=gym_config) diff --git a/embodichain/toolkits/scaffold/templates/extension/task_config_only.py.j2 b/embodichain/toolkits/scaffold/templates/extension/task_config_only.py.j2 new file mode 100644 index 00000000..642eb17e --- /dev/null +++ b/embodichain/toolkits/scaffold/templates/extension/task_config_only.py.j2 @@ -0,0 +1,28 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +from __future__ import annotations + +from embodichain.lab.gym.utils.registration import register_env +from embodichain.lab.gym.envs import EmbodiedEnv, EmbodiedEnvCfg + + +@register_env("{{ gym_id }}", max_episode_steps={{ max_episode_steps }}) +class {{ task_class }}(EmbodiedEnv): + """{{ description }}""" + + def __init__(self, cfg: EmbodiedEnvCfg = None, **kwargs): + super().__init__(cfg, **kwargs) diff --git a/embodichain/toolkits/scaffold/templates/extension/task_demo.py.j2 b/embodichain/toolkits/scaffold/templates/extension/task_demo.py.j2 new file mode 100644 index 00000000..f0380294 --- /dev/null +++ b/embodichain/toolkits/scaffold/templates/extension/task_demo.py.j2 @@ -0,0 +1,40 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +from __future__ import annotations + +import torch +from typing import Sequence + +from embodichain.lab.sim.types import EnvAction +from embodichain.lab.gym.utils.registration import register_env +from embodichain.lab.gym.envs import EmbodiedEnv, EmbodiedEnvCfg + + +@register_env("{{ gym_id }}", max_episode_steps={{ max_episode_steps }}) +class {{ task_class }}(EmbodiedEnv): + """{{ description }}""" + + def __init__(self, cfg: EmbodiedEnvCfg = None, **kwargs): + super().__init__(cfg, **kwargs) + + def create_demo_action_list(self, *args, **kwargs) -> Sequence[EnvAction] | None: + """Generate a demonstration action sequence for data collection.""" + demo_actions: list[EnvAction] = [] + for _ in range(50): + action = self.action_space.sample() * 0.05 + demo_actions.append(torch.as_tensor(action, device=self.device)) + return demo_actions diff --git a/embodichain/toolkits/scaffold/templates/extension/task_rl.py.j2 b/embodichain/toolkits/scaffold/templates/extension/task_rl.py.j2 new file mode 100644 index 00000000..ada74db5 --- /dev/null +++ b/embodichain/toolkits/scaffold/templates/extension/task_rl.py.j2 @@ -0,0 +1,44 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +from __future__ import annotations + +import torch +from typing import Any, Dict, Tuple + +from embodichain.lab.gym.utils.registration import register_env +from embodichain.lab.gym.envs import EmbodiedEnv, EmbodiedEnvCfg +from embodichain.lab.sim.types import EnvObs + + +@register_env("{{ gym_id }}", max_episode_steps={{ max_episode_steps }}, override=True) +class {{ task_class }}(EmbodiedEnv): + """{{ description }}""" + + def __init__(self, cfg: EmbodiedEnvCfg = None, **kwargs): + if cfg is None: + cfg = EmbodiedEnvCfg() + super().__init__(cfg, **kwargs) + + def compute_task_state( + self, **kwargs + ) -> Tuple[torch.Tensor, torch.Tensor, Dict[str, Any]]: + is_success = torch.zeros(self.num_envs, device=self.device, dtype=torch.bool) + is_fail = torch.zeros(self.num_envs, device=self.device, dtype=torch.bool) + return is_success, is_fail, {} + + def check_truncated(self, obs: EnvObs, info: Dict[str, Any]) -> torch.Tensor: + return torch.zeros(self.num_envs, device=self.device, dtype=torch.bool) diff --git a/embodichain/toolkits/scaffold/templates/extension/tasks_init.py.j2 b/embodichain/toolkits/scaffold/templates/extension/tasks_init.py.j2 new file mode 100644 index 00000000..35c9c6f2 --- /dev/null +++ b/embodichain/toolkits/scaffold/templates/extension/tasks_init.py.j2 @@ -0,0 +1,19 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +from __future__ import annotations + +from .{{ task_snake }} import {{ task_class }} diff --git a/embodichain/toolkits/scaffold/templates/extension/test_task.py.j2 b/embodichain/toolkits/scaffold/templates/extension/test_task.py.j2 new file mode 100644 index 00000000..713e0d46 --- /dev/null +++ b/embodichain/toolkits/scaffold/templates/extension/test_task.py.j2 @@ -0,0 +1,24 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +from __future__ import annotations + +import importlib + + +def test_package_import_registers_env(): + pkg = importlib.import_module("{{ package_name }}") + assert hasattr(pkg, "__version__") diff --git a/embodichain/toolkits/scaffold/templates/inrepo/task_config_only.py.j2 b/embodichain/toolkits/scaffold/templates/inrepo/task_config_only.py.j2 new file mode 100644 index 00000000..a48c969b --- /dev/null +++ b/embodichain/toolkits/scaffold/templates/inrepo/task_config_only.py.j2 @@ -0,0 +1,33 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +from __future__ import annotations + +from embodichain.lab.gym.utils.registration import register_env +from embodichain.lab.gym.envs import EmbodiedEnv, EmbodiedEnvCfg + +__all__ = ["{{ task_class }}"] + + +@register_env("{{ gym_id }}", max_episode_steps={{ max_episode_steps }}) +class {{ task_class }}(EmbodiedEnv): + """{{ description }} + + Scene, observations, and events are defined in the gym JSON config. + """ + + def __init__(self, cfg: EmbodiedEnvCfg = None, **kwargs): + super().__init__(cfg, **kwargs) diff --git a/embodichain/toolkits/scaffold/templates/inrepo/task_demo.py.j2 b/embodichain/toolkits/scaffold/templates/inrepo/task_demo.py.j2 new file mode 100644 index 00000000..54b3c55d --- /dev/null +++ b/embodichain/toolkits/scaffold/templates/inrepo/task_demo.py.j2 @@ -0,0 +1,45 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +from __future__ import annotations + +import torch +from typing import Sequence + +from embodichain.lab.sim.types import EnvAction +from embodichain.lab.gym.utils.registration import register_env +from embodichain.lab.gym.envs import EmbodiedEnv, EmbodiedEnvCfg + +__all__ = ["{{ task_class }}"] + + +@register_env("{{ gym_id }}", max_episode_steps={{ max_episode_steps }}) +class {{ task_class }}(EmbodiedEnv): + """{{ description }}""" + + def __init__(self, cfg: EmbodiedEnvCfg = None, **kwargs): + super().__init__(cfg, **kwargs) + + def create_demo_action_list(self, *args, **kwargs) -> Sequence[EnvAction] | None: + """Generate a demonstration action sequence for data collection. + + TODO: Replace random sampling with task-specific expert trajectories. + """ + demo_actions: list[EnvAction] = [] + for _ in range(50): + action = self.action_space.sample() * 0.05 + demo_actions.append(torch.as_tensor(action, device=self.device)) + return demo_actions diff --git a/embodichain/toolkits/scaffold/templates/inrepo/task_rl.py.j2 b/embodichain/toolkits/scaffold/templates/inrepo/task_rl.py.j2 new file mode 100644 index 00000000..f335af45 --- /dev/null +++ b/embodichain/toolkits/scaffold/templates/inrepo/task_rl.py.j2 @@ -0,0 +1,55 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +from __future__ import annotations + +import torch +from typing import Any, Dict, Tuple + +from embodichain.lab.gym.utils.registration import register_env +from embodichain.lab.gym.envs import EmbodiedEnv, EmbodiedEnvCfg +from embodichain.lab.sim.types import EnvObs + +__all__ = ["{{ task_class }}"] + + +@register_env("{{ gym_id }}", max_episode_steps={{ max_episode_steps }}, override=True) +class {{ task_class }}(EmbodiedEnv): + """{{ description }}""" + + def __init__(self, cfg: EmbodiedEnvCfg = None, **kwargs): + if cfg is None: + cfg = EmbodiedEnvCfg() + super().__init__(cfg, **kwargs) + + def compute_task_state( + self, **kwargs + ) -> Tuple[torch.Tensor, torch.Tensor, Dict[str, Any]]: + """Evaluate task success and failure per environment instance.""" + is_success = torch.zeros(self.num_envs, device=self.device, dtype=torch.bool) + is_fail = torch.zeros(self.num_envs, device=self.device, dtype=torch.bool) + metrics: Dict[str, Any] = {} + return is_success, is_fail, metrics + + def check_truncated(self, obs: EnvObs, info: Dict[str, Any]) -> torch.Tensor: + """Early episode truncation (e.g. object fallen).""" + return torch.zeros(self.num_envs, device=self.device, dtype=torch.bool) + +{% if spec.reward_style == "python" %} + def get_reward(self, obs: EnvObs, action, info: Dict[str, Any]) -> torch.Tensor: + """Per-step reward when not using JSON reward functors.""" + return torch.zeros(self.num_envs, device=self.device) +{% endif %} diff --git a/embodichain/toolkits/scaffold/templates/inrepo/test_task.py.j2 b/embodichain/toolkits/scaffold/templates/inrepo/test_task.py.j2 new file mode 100644 index 00000000..0e118be9 --- /dev/null +++ b/embodichain/toolkits/scaffold/templates/inrepo/test_task.py.j2 @@ -0,0 +1,36 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +from __future__ import annotations + +import pytest + +from embodichain.lab.gym.utils import registration as env_registry + +# Import registers the environment. +from embodichain.lab.gym.envs.tasks.{{ spec.category }}.{{ task_snake }} import ( # noqa: F401 + {{ task_class }}, +) + + +def test_env_registered(): + assert "{{ gym_id }}" in env_registry.REGISTERED_ENVS + + +@pytest.mark.gpu +def test_env_import_smoke(): + spec = env_registry.REGISTERED_ENVS["{{ gym_id }}"] + assert spec.cls.__name__ == "{{ task_class }}" diff --git a/pyproject.toml b/pyproject.toml index de1e5deb..451c6a0c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -48,9 +48,13 @@ dependencies = [ "h5py", "tensordict", "viser==1.0.21", - "lerobot>=0.4.4" + "lerobot>=0.4.4", + "jinja2>=3.1.0" ] +[project.scripts] +embodichain-new-task = "embodichain.toolkits.scaffold.cli:main" + [project.optional-dependencies] gensim = [ "bpy", @@ -71,4 +75,7 @@ version = { file = ["VERSION"] } where = ["."] exclude = ["docs"] +[tool.setuptools.package-data] +"embodichain.toolkits.scaffold" = ["templates/**/*.j2"] + [tool.black] diff --git a/setup.py b/setup.py index 55473076..d6646183 100644 --- a/setup.py +++ b/setup.py @@ -122,7 +122,14 @@ def main(): description="An end-to-end, GPU-accelerated, and modular platform for building generalized Embodied Intelligence.", packages=find_packages(exclude=["docs"]), data_files=data_files, - entry_points={}, + entry_points={ + "console_scripts": [ + "embodichain-new-task=embodichain.toolkits.scaffold.cli:main", + ], + }, + package_data={ + "embodichain.toolkits.scaffold": ["templates/**/*.j2"], + }, cmdclass=cmdclass, include_package_data=True, ) diff --git a/skills/add-task-env/SKILL.md b/skills/add-task-env/SKILL.md index b6092cfc..39c0605f 100644 --- a/skills/add-task-env/SKILL.md +++ b/skills/add-task-env/SKILL.md @@ -14,6 +14,30 @@ Scaffold a new task environment following EmbodiChain's conventions and patterns ## Steps +### 0. Scaffold with CLI (preferred) + +```bash +# In-repo expert demonstration task +embodichain-new-task \ + --target inrepo \ + --workflow demo \ + --category tableware \ + --name pick_place \ + --gym-id PickPlace-v1 + +# In-repo RL task +embodichain-new-task --target inrepo --workflow rl --name reach_cube --gym-id ReachCubeRL + +# External extension (see https://github.com/DexForce/embodichain_task_template) +embodichain-new-task --target extension --workflow demo --name my_task \ + --package-name my_tasks --output-dir ./my_tasks + +# Interactive prompts +embodichain-new-task --interactive +``` + +Then customize generated files (scene JSON, `create_demo_action_list`, RL hooks). + ### 1. Determine Task Category Ask the user: diff --git a/tests/toolkits/scaffold/test_scaffold.py b/tests/toolkits/scaffold/test_scaffold.py new file mode 100644 index 00000000..6cc07f70 --- /dev/null +++ b/tests/toolkits/scaffold/test_scaffold.py @@ -0,0 +1,115 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +from __future__ import annotations + +import json +import re +from pathlib import Path + +import pytest + +from embodichain.toolkits.scaffold.generator import generate_task +from embodichain.toolkits.scaffold.naming import default_env_class, default_gym_id +from embodichain.toolkits.scaffold.post_process import patch_tasks_init +from embodichain.toolkits.scaffold.render import render_task_py +from embodichain.toolkits.scaffold.spec import TaskSpec + + +def test_naming_defaults(): + assert default_env_class("pick_place") == "PickPlaceEnv" + assert default_gym_id("pick_place") == "PickPlace-v1" + + +def test_invalid_snake_raises(): + with pytest.raises(ValueError, match="snake_case"): + TaskSpec(task_snake="PickPlace", workflow="demo") + + +def test_render_demo_contains_register(): + spec = TaskSpec( + task_snake="scaffold_test_demo", + workflow="demo", + target="inrepo", + category="special", + gym_id="ScaffoldTestDemo-v0", + task_class="ScaffoldTestDemoEnv", + include_test=False, + ) + text = render_task_py(spec) + assert '@register_env("ScaffoldTestDemo-v0"' in text + assert "create_demo_action_list" in text + + +def test_dry_run_inrepo_paths(tmp_path: Path): + spec = TaskSpec( + task_snake="dry_run_task", + workflow="config-only", + target="inrepo", + category="special", + gym_id="DryRunTask-v0", + dry_run=True, + include_test=True, + repo_root=tmp_path, + ) + paths = generate_task(spec) + assert not spec.task_py_path().exists() + assert len(paths) >= 2 + + +def test_generate_extension_project(tmp_path: Path): + out = tmp_path / "my_ext" + spec = TaskSpec( + task_snake="ext_task", + workflow="demo", + target="extension", + gym_id="ExtTask-v1", + package_name="my_ext_pkg", + project_name="my-ext", + output_dir=out, + include_test=True, + run_black=False, + repo_root=tmp_path, + ) + written = generate_task(spec) + assert (out / "pyproject.toml").is_file() + assert (out / "my_ext_pkg" / "tasks" / "ext_task.py").is_file() + assert (out / "configs" / "ext_task" / "gym.json").is_file() + gym = json.loads((out / "configs" / "ext_task" / "gym.json").read_text()) + assert gym["id"] == "ExtTask-v1" + init = (out / "my_ext_pkg" / "tasks" / "__init__.py").read_text() + assert "from .ext_task import ExtTaskEnv" in init + assert len(written) >= 5 + + +def test_patch_tasks_init_idempotent(tmp_path: Path): + spec = TaskSpec( + task_snake="foo", + workflow="demo", + target="inrepo", + category="special", + gym_id="Foo-v1", + task_class="FooEnv", + repo_root=tmp_path, + ) + init = spec.tasks_init_path() + init.parent.mkdir(parents=True, exist_ok=True) + init.write_text("__all__ = []\n", encoding="utf-8") + patch_tasks_init(spec) + patch_tasks_init(spec) + content = init.read_text() + assert "FooEnv" in content + assert "from embodichain.lab.gym.envs.tasks.special.foo import FooEnv" in content