From d75a5b8f0ac4a5759c2838f158962fa6577cabe5 Mon Sep 17 00:00:00 2001 From: yuecideng Date: Mon, 6 Jul 2026 17:26:20 +0000 Subject: [PATCH 1/2] docs: add task environment decoupling design spec Co-Authored-By: Claude --- .../2026-07-07-task-env-refactor-design.md | 458 ++++++++++++++++++ 1 file changed, 458 insertions(+) create mode 100644 docs/superpowers/specs/2026-07-07-task-env-refactor-design.md diff --git a/docs/superpowers/specs/2026-07-07-task-env-refactor-design.md b/docs/superpowers/specs/2026-07-07-task-env-refactor-design.md new file mode 100644 index 00000000..2f9e64ba --- /dev/null +++ b/docs/superpowers/specs/2026-07-07-task-env-refactor-design.md @@ -0,0 +1,458 @@ +# Task Environment Decoupling: Design Spec + +**Date**: 2026-07-07 +**Status**: Draft +**Reference**: IsaacLab task architecture analysis + +## 1. Motivation + +Currently all task environments live inside the core `embodichain` package under +`embodichain/lab/gym/envs/tasks/`. This creates several problems: + +- **Core package bloat** — tableware, RL, and special tasks are bundled with + framework code, increasing install size and cognitive load. +- **No third-party extension path** — downstream projects like RoboSynChallenge + must monkey-patch internals (`DEFAULT_MANAGER_MODULES`, `get_data_path`, + `LeRobotRecorder`) to integrate. +- **Implicit dependency** — downstream projects use `sys.path` hacks instead of + proper Python package dependencies. +- **Duplicated launch scripts** — every downstream project copies its own + `run_env.py` from the core. + +## 2. Goals + +1. **Extract all task environments** from `embodichain` into a separate + `embodichain_tasks` package within the same repository. +2. **Enable third-party task packages** via standard setuptools + `entry_points`, allowing external projects to register tasks and init hooks + without modifying `embodichain` source. +3. **Provide a unified launch script** — one `run_env` CLI that discovers and + launches any registered task from any installed package. +4. **Full backward compatibility** — existing env IDs, JSON/YAML config formats, + `@register_env` API, and `gym.make()` workflow remain unchanged. + +## 3. Non-Goals (this iteration) + +- RoboSynChallenge is NOT modified in this iteration. It will adopt the new + mechanism in a follow-up. +- `BaseAgentEnv` is moved out with tasks but will be replaced by a new Agent + base class in future work. No redesign of the agent infrastructure. +- Config format remains JSON/YAML. No migration to Python config classes. + +## 4. Architecture Overview + +``` +┌──────────────────────────────────────────────────────┐ +│ python -m embodichain.lab.scripts.run_env │ +│ (unified CLI, no task-specific imports) │ +└──────────┬───────────────────────────────────────────┘ + │ + 1. discover_task_packages() + └─ importlib.metadata.entry_points("embodichain.tasks") + imports: embodichain_tasks, robosynchallenge, ... + └─ triggers @register_env → gym.register() + │ + 2. execute_init_hooks() + └─ importlib.metadata.entry_points("embodichain.init") + calls: robosynchallenge.init:register() + └─ register_manager_modules(...) + └─ install_asset_resolver() + └─ install_lerobot_recorder_override() + │ + 3. build_env_cfg_from_args(args) + └─ loads gym_config.json → EmbodiedEnvCfg + │ + 4. gym.make(id, cfg=env_cfg) + └─ gymnasium registry → instantiated env +``` + +## 5. Package Structure Changes + +### 5.1 New package: `embodichain_tasks/` + +``` +embodichain_tasks/ +├── pyproject.toml +├── VERSION +├── embodichain_tasks/ +│ ├── __init__.py # import_packages() recursive import +│ ├── utils/ +│ │ └── importer.py # import_packages() utility +│ ├── tableware/ +│ │ ├── __init__.py +│ │ ├── base_agent_env.py # migrated from core (to be deprecated) +│ │ ├── pour_water/ +│ │ │ ├── __init__.py +│ │ │ ├── pour_water.py # @register_env("PourWater-v3") +│ │ │ └── action_bank.py +│ │ ├── rearrangement.py +│ │ ├── stack_blocks_two.py +│ │ ├── stack_cups.py +│ │ ├── scoop_ice.py +│ │ ├── blocks_ranking_rgb.py +│ │ ├── blocks_ranking_size.py +│ │ ├── match_object_container.py +│ │ └── place_object_drawer.py +│ ├── rl/ +│ │ ├── __init__.py +│ │ ├── build_env.py # migrated from tasks/rl/__init__.py +│ │ ├── push_cube.py +│ │ └── basic/ +│ │ └── cart_pole.py +│ └── special/ +│ └── simple_task.py +└── configs/ # task JSON configs (migrated from various locations) + ├── pour_water/ + ├── rearrangement/ + └── ... +``` + +### 5.2 Core package `embodichain/` removals + +``` +embodichain/lab/gym/envs/tasks/ # ENTIRE DIRECTORY REMOVED +``` + +These stay in core (framework-level, not task-specific): +- `embodichain/lab/gym/envs/base_env.py` — `BaseEnv`, `EnvCfg` +- `embodichain/lab/gym/envs/embodied_env.py` — `EmbodiedEnv`, `EmbodiedEnvCfg` +- `embodichain/lab/gym/envs/managers/` — all manager infrastructure +- `embodichain/lab/gym/envs/action_bank/` — `ActionBank`, `@tag_node`, `@tag_edge` +- `embodichain/lab/gym/envs/wrapper/` — `NoFailWrapper` +- `embodichain/lab/gym/utils/` — `gym_utils`, `registration`, `misc` + +### 5.3 Core package additions + +- `embodichain/lab/gym/utils/registration.py` — new functions: + - `discover_task_packages()` — entry_points-based task discovery + - `execute_init_hooks()` — entry_points-based init hook execution +- `embodichain/lab/gym/utils/gym_utils.py` — new functions: + - `register_manager_modules(modules)` — register custom manager modules + - `get_manager_modules()` — get all registered manager modules + +## 6. Registration & Discovery Mechanism + +### 6.1 setuptools entry_points + +Task packages declare themselves in `pyproject.toml`: + +```toml +# embodichain_tasks/pyproject.toml +[project.entry-points."embodichain.tasks"] +"embodichain_tasks" = "embodichain_tasks" +``` + +Third-party packages follow the same pattern: + +```toml +# robosynchallenge/pyproject.toml (future) +[project.entry-points."embodichain.tasks"] +"robosynchallenge" = "robosynchallenge.tasks" +``` + +### 6.2 Discovery function + +```python +# embodichain/lab/gym/utils/registration.py + +import importlib.metadata + +def discover_task_packages() -> list[str]: + """Import all registered task packages via entry_points. + + Returns: + List of imported package names. + """ + imported = [] + for ep in importlib.metadata.entry_points(group="embodichain.tasks"): + importlib.import_module(ep.value) + imported.append(ep.name) + return imported +``` + +### 6.3 Task package auto-import + +Each task package's `__init__.py` uses `import_packages()` (same pattern as +IsaacLab's `isaaclab_tasks`): + +```python +# embodichain_tasks/__init__.py +from .utils.importer import import_packages +_BLACKLIST = ["utils"] +import_packages(__name__, _BLACKLIST) +``` + +`import_packages()` recursively walks all sub-packages, importing each module. +This triggers each task's `__init__.py`, which calls `@register_env` → +`gym.register()`. After `discover_task_packages()` returns, all tasks from all +installed packages are registered in `gymnasium.registry`. + +### 6.4 @register_env — unchanged + +The `@register_env` decorator continues to: +1. Record the env class in `REGISTERED_ENVS` dict +2. Call `gym.register()` with `entry_point=partial(make, env_id=uid)` + +All existing env IDs remain unchanged (e.g., `PourWater-v3`, `PushCubeRL`, +`CartPoleRL`, `StackBlocksTwo-v1`, etc.). + +## 7. Init Hooks Mechanism + +### 7.1 Declaration + +Task packages optionally declare init hooks: + +```toml +[project.entry-points."embodichain.init"] +"robosynchallenge" = "robosynchallenge.init:register" +``` + +### 7.2 Hook signature + +```python +# robosynchallenge/init.py (future) + +def register() -> None: + """EmbodiChain init hook — called before env creation.""" + from embodichain.lab.gym.utils.gym_utils import register_manager_modules + from robosynchallenge.data.asset_resolver import install_embodichain_asset_resolver + from robosynchallenge.managers.datasets import install_lerobot_recorder_override + + register_manager_modules([ + "robosynchallenge.managers.actions", + "robosynchallenge.managers.datasets", + "robosynchallenge.managers.events", + "robosynchallenge.managers.observations", + ]) + install_embodichain_asset_resolver() + install_lerobot_recorder_override() +``` + +### 7.3 Hook execution + +```python +# embodichain/lab/gym/utils/registration.py + +def execute_init_hooks() -> list[str]: + """Execute all registered init hooks. + + Hooks are called in entry_points declaration order. + Exceptions from one hook do not prevent others from executing. + + Returns: + List of hook names that were executed successfully. + """ + executed = [] + for ep in importlib.metadata.entry_points(group="embodichain.init"): + try: + module_name, func_name = ep.value.split(":") + module = importlib.import_module(module_name) + func = getattr(module, func_name) + func() + executed.append(ep.name) + except Exception: + import logging + logging.getLogger(__name__).warning( + f"Init hook '{ep.name}' failed", exc_info=True + ) + return executed +``` + +### 7.4 manager module registration API + +```python +# embodichain/lab/gym/utils/gym_utils.py + +_EXTRA_MANAGER_MODULES: list[str] = [] + +def register_manager_modules(modules: list[str]) -> None: + """Register additional manager modules for functor resolution. + + These modules are searched when resolving functor function names + from config strings (in addition to the built-in DEFAULT_MANAGER_MODULES). + """ + for m in modules: + if m not in _EXTRA_MANAGER_MODULES: + _EXTRA_MANAGER_MODULES.append(m) + +def get_manager_modules() -> list[str]: + """Get all registered manager modules.""" + return DEFAULT_MANAGER_MODULES + _EXTRA_MANAGER_MODULES +``` + +All call sites that currently reference `DEFAULT_MANAGER_MODULES` directly +(e.g., `config_to_cfg()`, `find_function_from_modules()`) are updated to use +`get_manager_modules()`. + +## 8. Unified Launch Script + +### 8.1 Changes to `run_env.py` + +```python +# embodichain/lab/scripts/run_env.py + +def main(): + args = parse_args() + + # Step 1: Discover all task packages via entry_points + discover_task_packages() + + # Step 2: Execute init hooks (register managers, asset resolvers, etc.) + execute_init_hooks() + + # Step 3: Build env config from CLI args + JSON config file + env_cfg = build_env_cfg_from_args(args) + + # Step 4: Create environment via gymnasium + env = gym.make(id=env_cfg["id"], cfg=env_cfg) + + # Step 5: Run data generation or preview + if args.preview: + preview(env) + else: + generate_and_execute_action_list(env, ...) +``` + +### 8.2 CLI — unchanged + +All existing CLI arguments remain: + +``` +--gym_config Path to gym config JSON/YAML +--action_config Path to action config JSON +--num_envs Number of parallel environments +--device cpu | cuda +--headless Run headless +--renderer auto | hybrid | fast-rt | rt +--arena_space Arena space size +--gpu_id GPU device ID +--preview Interactive preview mode +--filter_visual_rand Disable visual randomization +--filter_dataset_saving Disable dataset saving +--max_episodes Override max episodes +``` + +### 8.3 Cross-package config paths + +Config files can live anywhere on the filesystem — the `run_env` script does +not assume configs are inside the `embodichain` package: + +```bash +# Official tasks +python -m embodichain.lab.scripts.run_env \ + --gym_config ../embodichain_tasks/configs/pour_water/random/gym_config.json + +# Third-party tasks (future) +python -m embodichain.lab.scripts.run_env \ + --gym_config ../RoboSynChallenge/configs/click_bell/random/gym_config.json +``` + +Task identity is determined solely by the `"id"` field in the gym config JSON, +which must match a registered gym environment ID. + +## 9. Backward Compatibility + +### 9.1 Guarantees + +| Item | Strategy | +|------|----------| +| Env IDs | Unchanged — `@register_env(uid)` values preserved | +| JSON/YAML config format | Unchanged — same schema | +| `gym.make(id, cfg=...)` | Unchanged — standard gymnasium interface | +| `@register_env` API | Unchanged — same signature | +| `EmbodiedEnv` / `BaseEnv` import path | Unchanged — stays in `embodichain.lab.gym.envs` | +| `Functor` / manager infrastructure | Unchanged — stays in `embodichain.lab.gym.envs.managers` | +| `ActionBank` | Unchanged — stays in `embodichain.lab.gym.envs.action_bank` | +| `build_env_cfg_from_args()` / `config_to_cfg()` | Unchanged — stays in `embodichain.lab.gym.utils` | +| `run_env` CLI arguments | Unchanged — same flags and semantics | + +### 9.2 Breaking change + +Direct imports of task classes from `embodichain.lab.gym.envs.tasks.*` will +break. This is not the standard usage path (tasks are created via +`gym.make(env_id)`), but any code doing direct imports needs to update the +import path to `embodichain_tasks.*`. + +### 9.3 Deprecation shim + +`embodichain/lab/gym/envs/tasks/__init__.py` becomes a thin shim: + +```python +# This module is deprecated. Import from embodichain_tasks instead. +import warnings +warnings.warn( + "embodichain.lab.gym.envs.tasks is deprecated. " + "Import from embodichain_tasks instead.", + DeprecationWarning, + stacklevel=2, +) + +# Re-export everything from the new package for backward compatibility +from embodichain_tasks.tableware.base_agent_env import BaseAgentEnv # noqa +from embodichain_tasks.tableware.pour_water.pour_water import ( + PourWaterEnv, PourWaterAgentEnv, +) +from embodichain_tasks.tableware.rearrangement import ( + RearrangementEnv, RearrangementAgentEnv, +) +# ... (all other task classes) +``` + +This ensures existing direct imports continue to work with a deprecation +warning. The shim is removed in the next major version. + +## 12. Resolved Decisions + +- **Config location**: `embodichain_tasks/configs/` — configs ship with the + task package, not in a separate top-level directory. +- **Package type**: `embodichain_tasks` is a regular pip-installable package + that depends on `embodichain`. Not a namespace package. +- **`import_packages()` utility**: A lightweight implementation lives in + `embodichain_tasks/utils/importer.py`, following the IsaacLab pattern but + without Omniverse-specific logic. +- **Entry point group names**: `embodichain.tasks` for task packages, + `embodichain.init` for init hooks. Both follow the convention + `.`. +- **Hook execution order**: Entry points are iterated in declaration order. + No priority system — if ordering matters, the downstream package controls + it by its position in the entry_points list. + +## 10. Third-Party Integration Pattern (Reference) + +When RoboSynChallenge (or any downstream project) adopts this mechanism, the +changes are minimal: + +### 10.1 Changes required (3 files) + +| File | Change | +|------|--------| +| `pyproject.toml` | Add `[project.entry-points."embodichain.tasks"]` and `[project.entry-points."embodichain.init"]` | +| `robosynchallenge/init.py` | New file with `register()` function | +| `scripts/run_env.py` | Delete (use core `run_env`) | + +### 10.2 No changes required + +- All JSON config files +- All `@register_env` decorators +- All task environment classes +- All `ActionBank` subclasses +- All custom functors +- Shell launch scripts (just update the `RUN_ENV` variable) + +## 11. Implementation Plan (High-Level) + +1. Create `embodichain_tasks/` package skeleton with `pyproject.toml` +2. Port `import_packages()` utility from IsaacLab pattern +3. Migrate all task modules from `embodichain/lab/gym/envs/tasks/` to + `embodichain_tasks/embodichain_tasks/`, updating internal imports +4. Migrate task JSON configs to `embodichain_tasks/configs/` +5. Add `discover_task_packages()` and `execute_init_hooks()` to + `embodichain/lab/gym/utils/registration.py` +6. Add `register_manager_modules()` / `get_manager_modules()` to + `embodichain/lab/gym/utils/gym_utils.py` +7. Update all `DEFAULT_MANAGER_MODULES` references to use `get_manager_modules()` +8. Update `run_env.py` to call `discover_task_packages()` + `execute_init_hooks()` +9. Add deprecation shim in `embodichain/lab/gym/envs/tasks/__init__.py` +10. Run full test suite; verify all tasks launch via unified `run_env` +11. Update project documentation + From 9e6c8d9b8e0a7d6fa693def841c1d60804619839 Mon Sep 17 00:00:00 2001 From: yuecideng Date: Mon, 6 Jul 2026 18:07:03 +0000 Subject: [PATCH 2/2] refactor: extract task environments into separate embodichain_tasks package - Create embodichain_tasks/ as a separate pip-installable package - Migrate all task envs from embodichain/lab/gym/envs/tasks/ to embodichain_tasks/embodichain_tasks/ - Add entry_points-based task discovery (embodichain.tasks group) - Add init hooks mechanism (embodichain.init group) for third-party manager modules, asset resolvers, etc. - Add register_manager_modules() / get_manager_modules() to gym_utils - Update run_env.py to auto-discover tasks via entry_points - Add deprecation shim in tasks/__init__.py for backward compatibility - All existing env IDs, config formats, and APIs unchanged Co-Authored-By: Claude --- .../plans/2026-07-07-task-env-refactor.md | 19 + embodichain/lab/gym/envs/__init__.py | 9 +- embodichain/lab/gym/envs/tasks/__init__.py | 115 ++- embodichain/lab/gym/utils/gym_utils.py | 33 +- embodichain/lab/gym/utils/registration.py | 69 ++ embodichain/lab/scripts/run_env.py | 10 + embodichain_tasks/VERSION | 1 + .../agents/rl/basic/cart_pole/gym_config.json | 74 ++ .../agents/rl/basic/cart_pole/gym_config.yaml | 66 ++ .../rl/basic/cart_pole/train_config.json | 93 ++ .../rl/basic/cart_pole/train_config.yaml | 71 ++ .../rl/basic/cart_pole/train_config_grpo.json | 87 ++ .../rl/basic/cart_pole/train_config_grpo.yaml | 68 ++ .../agents/rl/push_cube/gym_config.json | 185 ++++ .../agents/rl/push_cube/train_config.json | 76 ++ .../rl/push_cube/train_config_grpo.json | 65 ++ .../configs/gym/action_bank/conf.json | 237 +++++ .../agent/pour_water_agent/agent_config.json | 31 + .../pour_water_agent/agent_config_dual.json | 31 + .../pour_water_agent/fast_gym_config.json | 391 ++++++++ .../agent/pour_water_agent/task_prompt.txt | 5 + .../pour_water_agent/task_prompt_dual.txt | 5 + .../rearrangement_agent/agent_config.json | 31 + .../rearrangement_agent/fast_gym_config.json | 383 +++++++ .../agent/rearrangement_agent/task_prompt.txt | 8 + .../blocks_ranking_rgb/cobot_magic_3cam.json | 283 ++++++ .../blocks_ranking_size/cobot_magic_3cam.json | 257 +++++ embodichain_tasks/configs/gym/cobotmagic.json | 121 +++ embodichain_tasks/configs/gym/cobotmagic.yaml | 153 +++ .../configs/gym/dexforce_w1.json | 109 ++ .../configs/gym/dexforce_w1.yaml | 134 +++ .../cobot_magic_3cam.json | 343 +++++++ .../place_object_drawer/cobot_magic_3cam.json | 165 +++ .../configs/gym/pour_water/action_config.json | 938 ++++++++++++++++++ .../configs/gym/pour_water/gym_config.json | 351 +++++++ .../configs/gym/scoop_ice/gym_config.json | 235 +++++ .../configs/gym/special/simple_task_ur10.json | 104 ++ .../stack_blocks_two/cobot_magic_3cam.json | 191 ++++ .../gym/stack_cups/cobot_magic_3cam.json | 202 ++++ .../embodichain_tasks/__init__.py | 31 + .../embodichain_tasks/rl/__init__.py | 32 + .../embodichain_tasks/rl/basic/__init__.py | 0 .../embodichain_tasks/rl/basic/cart_pole.py | 78 ++ .../embodichain_tasks/rl/push_cube.py | 65 ++ .../embodichain_tasks/special/__init__.py | 0 .../embodichain_tasks/special/simple_task.py | 87 ++ .../embodichain_tasks/tableware/__init__.py | 0 .../tableware/base_agent_env.py | 201 ++++ .../tableware/blocks_ranking_rgb.py | 324 ++++++ .../tableware/blocks_ranking_size.py | 89 ++ .../tableware/match_object_container.py | 176 ++++ .../tableware/place_object_drawer.py | 83 ++ .../tableware/pour_water/__init__.py | 0 .../tableware/pour_water/action_bank.py | 255 +++++ .../tableware/pour_water/pour_water.py | 161 +++ .../tableware/rearrangement.py | 90 ++ .../embodichain_tasks/tableware/scoop_ice.py | 213 ++++ .../tableware/stack_blocks_two.py | 324 ++++++ .../embodichain_tasks/tableware/stack_cups.py | 118 +++ .../embodichain_tasks/utils/__init__.py | 17 + .../embodichain_tasks/utils/importer.py | 52 + embodichain_tasks/pyproject.toml | 18 + 62 files changed, 8112 insertions(+), 51 deletions(-) create mode 100644 docs/superpowers/plans/2026-07-07-task-env-refactor.md create mode 100644 embodichain_tasks/VERSION create mode 100644 embodichain_tasks/configs/agents/rl/basic/cart_pole/gym_config.json create mode 100644 embodichain_tasks/configs/agents/rl/basic/cart_pole/gym_config.yaml create mode 100644 embodichain_tasks/configs/agents/rl/basic/cart_pole/train_config.json create mode 100644 embodichain_tasks/configs/agents/rl/basic/cart_pole/train_config.yaml create mode 100644 embodichain_tasks/configs/agents/rl/basic/cart_pole/train_config_grpo.json create mode 100644 embodichain_tasks/configs/agents/rl/basic/cart_pole/train_config_grpo.yaml create mode 100644 embodichain_tasks/configs/agents/rl/push_cube/gym_config.json create mode 100644 embodichain_tasks/configs/agents/rl/push_cube/train_config.json create mode 100644 embodichain_tasks/configs/agents/rl/push_cube/train_config_grpo.json create mode 100644 embodichain_tasks/configs/gym/action_bank/conf.json create mode 100644 embodichain_tasks/configs/gym/agent/pour_water_agent/agent_config.json create mode 100644 embodichain_tasks/configs/gym/agent/pour_water_agent/agent_config_dual.json create mode 100644 embodichain_tasks/configs/gym/agent/pour_water_agent/fast_gym_config.json create mode 100644 embodichain_tasks/configs/gym/agent/pour_water_agent/task_prompt.txt create mode 100644 embodichain_tasks/configs/gym/agent/pour_water_agent/task_prompt_dual.txt create mode 100644 embodichain_tasks/configs/gym/agent/rearrangement_agent/agent_config.json create mode 100644 embodichain_tasks/configs/gym/agent/rearrangement_agent/fast_gym_config.json create mode 100644 embodichain_tasks/configs/gym/agent/rearrangement_agent/task_prompt.txt create mode 100644 embodichain_tasks/configs/gym/blocks_ranking_rgb/cobot_magic_3cam.json create mode 100644 embodichain_tasks/configs/gym/blocks_ranking_size/cobot_magic_3cam.json create mode 100644 embodichain_tasks/configs/gym/cobotmagic.json create mode 100644 embodichain_tasks/configs/gym/cobotmagic.yaml create mode 100644 embodichain_tasks/configs/gym/dexforce_w1.json create mode 100644 embodichain_tasks/configs/gym/dexforce_w1.yaml create mode 100644 embodichain_tasks/configs/gym/match_object_container/cobot_magic_3cam.json create mode 100644 embodichain_tasks/configs/gym/place_object_drawer/cobot_magic_3cam.json create mode 100644 embodichain_tasks/configs/gym/pour_water/action_config.json create mode 100644 embodichain_tasks/configs/gym/pour_water/gym_config.json create mode 100644 embodichain_tasks/configs/gym/scoop_ice/gym_config.json create mode 100644 embodichain_tasks/configs/gym/special/simple_task_ur10.json create mode 100644 embodichain_tasks/configs/gym/stack_blocks_two/cobot_magic_3cam.json create mode 100644 embodichain_tasks/configs/gym/stack_cups/cobot_magic_3cam.json create mode 100644 embodichain_tasks/embodichain_tasks/__init__.py create mode 100644 embodichain_tasks/embodichain_tasks/rl/__init__.py create mode 100644 embodichain_tasks/embodichain_tasks/rl/basic/__init__.py create mode 100644 embodichain_tasks/embodichain_tasks/rl/basic/cart_pole.py create mode 100644 embodichain_tasks/embodichain_tasks/rl/push_cube.py create mode 100644 embodichain_tasks/embodichain_tasks/special/__init__.py create mode 100644 embodichain_tasks/embodichain_tasks/special/simple_task.py create mode 100644 embodichain_tasks/embodichain_tasks/tableware/__init__.py create mode 100644 embodichain_tasks/embodichain_tasks/tableware/base_agent_env.py create mode 100644 embodichain_tasks/embodichain_tasks/tableware/blocks_ranking_rgb.py create mode 100644 embodichain_tasks/embodichain_tasks/tableware/blocks_ranking_size.py create mode 100644 embodichain_tasks/embodichain_tasks/tableware/match_object_container.py create mode 100644 embodichain_tasks/embodichain_tasks/tableware/place_object_drawer.py create mode 100644 embodichain_tasks/embodichain_tasks/tableware/pour_water/__init__.py create mode 100644 embodichain_tasks/embodichain_tasks/tableware/pour_water/action_bank.py create mode 100644 embodichain_tasks/embodichain_tasks/tableware/pour_water/pour_water.py create mode 100644 embodichain_tasks/embodichain_tasks/tableware/rearrangement.py create mode 100644 embodichain_tasks/embodichain_tasks/tableware/scoop_ice.py create mode 100644 embodichain_tasks/embodichain_tasks/tableware/stack_blocks_two.py create mode 100644 embodichain_tasks/embodichain_tasks/tableware/stack_cups.py create mode 100644 embodichain_tasks/embodichain_tasks/utils/__init__.py create mode 100644 embodichain_tasks/embodichain_tasks/utils/importer.py create mode 100644 embodichain_tasks/pyproject.toml diff --git a/docs/superpowers/plans/2026-07-07-task-env-refactor.md b/docs/superpowers/plans/2026-07-07-task-env-refactor.md new file mode 100644 index 00000000..bd779dde --- /dev/null +++ b/docs/superpowers/plans/2026-07-07-task-env-refactor.md @@ -0,0 +1,19 @@ +# Task Environment Decoupling — Implementation Plan + +> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. + +**Goal:** Extract all task envs from `embodichain` into `embodichain_tasks` package, add entry_points discovery, init hooks, and unified launch. + +**Architecture:** IsaacLab-style separation — `embodichain_tasks/` is a separate pip-installable package with its own `pyproject.toml`. Core `embodichain` discovers tasks via `importlib.metadata.entry_points(group="embodichain.tasks")`. + +**Tech Stack:** Python 3.10+, setuptools, gymnasium, importlib.metadata + +## Global Constraints + +- All existing env IDs unchanged +- JSON/YAML config format unchanged +- `@register_env` API unchanged +- `gym.make()` workflow unchanged +- RoboSynChallenge NOT touched this iteration + +--- diff --git a/embodichain/lab/gym/envs/__init__.py b/embodichain/lab/gym/envs/__init__.py index f7579fe4..d369a7b7 100644 --- a/embodichain/lab/gym/envs/__init__.py +++ b/embodichain/lab/gym/envs/__init__.py @@ -16,5 +16,12 @@ from .base_env import * from .embodied_env import * -from .tasks import * from .wrapper import * + +# Tasks have been moved to the ``embodichain_tasks`` package. +# The shim below preserves backward compatibility for direct imports +# from ``embodichain.lab.gym.envs.tasks``. +try: + from .tasks import * # noqa: F403 +except ImportError: + pass diff --git a/embodichain/lab/gym/envs/tasks/__init__.py b/embodichain/lab/gym/envs/tasks/__init__.py index 56da5a8c..48ad24d7 100644 --- a/embodichain/lab/gym/envs/tasks/__init__.py +++ b/embodichain/lab/gym/envs/tasks/__init__.py @@ -14,55 +14,76 @@ # limitations under the License. # ---------------------------------------------------------------------------- +"""Deprecation shim — task environments have moved to ``embodichain_tasks``. + +This module re-exports all task classes from the new ``embodichain_tasks`` +package for backward compatibility. It will be removed in a future version. + +.. deprecated:: + Import from ``embodichain_tasks`` directly instead of + ``embodichain.lab.gym.envs.tasks``. +""" + from __future__ import annotations -# Tableware task environments -from embodichain.lab.gym.envs.tasks.tableware.pour_water.pour_water import ( - PourWaterEnv, - PourWaterAgentEnv, -) -from embodichain.lab.gym.envs.tasks.tableware.scoop_ice import ScoopIce -from embodichain.lab.gym.envs.tasks.tableware.stack_blocks_two import StackBlocksTwoEnv -from embodichain.lab.gym.envs.tasks.tableware.blocks_ranking_rgb import ( - BlocksRankingRGBEnv, -) -from embodichain.lab.gym.envs.tasks.tableware.blocks_ranking_size import ( - BlocksRankingSizeEnv, -) -from embodichain.lab.gym.envs.tasks.tableware.place_object_drawer import ( - PlaceObjectDrawerEnv, -) -from embodichain.lab.gym.envs.tasks.tableware.stack_cups import ( - StackCupsEnv, -) -from embodichain.lab.gym.envs.tasks.tableware.match_object_container import ( - MatchObjectContainerEnv, -) -from embodichain.lab.gym.envs.tasks.tableware.rearrangement import ( - RearrangementEnv, - RearrangementAgentEnv, -) +import warnings -# Reinforcement learning environments -from embodichain.lab.gym.envs.tasks.rl.push_cube import PushCubeEnv -from embodichain.lab.gym.envs.tasks.rl.basic.cart_pole import CartPoleEnv +warnings.warn( + "embodichain.lab.gym.envs.tasks is deprecated. " + "Import from embodichain_tasks instead.", + DeprecationWarning, + stacklevel=2, +) -# Special environments -from embodichain.lab.gym.envs.tasks.special.simple_task import SimpleTaskEnv +try: + from embodichain_tasks.tableware.base_agent_env import BaseAgentEnv # noqa: F401 + from embodichain_tasks.tableware.pour_water.pour_water import ( # noqa: F401 + PourWaterEnv, + PourWaterAgentEnv, + ) + from embodichain_tasks.tableware.rearrangement import ( # noqa: F401 + RearrangementEnv, + RearrangementAgentEnv, + ) + from embodichain_tasks.tableware.stack_blocks_two import ( + StackBlocksTwoEnv, + ) # noqa: F401 + from embodichain_tasks.tableware.stack_cups import StackCupsEnv # noqa: F401 + from embodichain_tasks.tableware.scoop_ice import ScoopIce # noqa: F401 + from embodichain_tasks.tableware.blocks_ranking_rgb import ( # noqa: F401 + BlocksRankingRGBEnv, + ) + from embodichain_tasks.tableware.blocks_ranking_size import ( # noqa: F401 + BlocksRankingSizeEnv, + ) + from embodichain_tasks.tableware.match_object_container import ( # noqa: F401 + MatchObjectContainerEnv, + ) + from embodichain_tasks.tableware.place_object_drawer import ( # noqa: F401 + PlaceObjectDrawerEnv, + ) + from embodichain_tasks.rl.push_cube import PushCubeEnv # noqa: F401 + from embodichain_tasks.rl.basic.cart_pole import CartPoleEnv # noqa: F401 + from embodichain_tasks.special.simple_task import SimpleTaskEnv # noqa: F401 -__all__ = [ - "PourWaterEnv", - "PourWaterAgentEnv", - "ScoopIce", - "StackBlocksTwoEnv", - "BlocksRankingRGBEnv", - "BlocksRankingSizeEnv", - "PlaceObjectDrawerEnv", - "StackCupsEnv", - "MatchObjectContainerEnv", - "RearrangementEnv", - "RearrangementAgentEnv", - "PushCubeEnv", - "CartPoleEnv", - "SimpleTaskEnv", -] + __all__ = [ + "BaseAgentEnv", + "PourWaterEnv", + "PourWaterAgentEnv", + "RearrangementEnv", + "RearrangementAgentEnv", + "StackBlocksTwoEnv", + "StackCupsEnv", + "ScoopIce", + "BlocksRankingRGBEnv", + "BlocksRankingSizeEnv", + "MatchObjectContainerEnv", + "PlaceObjectDrawerEnv", + "PushCubeEnv", + "CartPoleEnv", + "SimpleTaskEnv", + ] +except ImportError: + # embodichain_tasks is not installed — tasks will be discovered via + # entry_points instead. + __all__: list[str] = [] diff --git a/embodichain/lab/gym/utils/gym_utils.py b/embodichain/lab/gym/utils/gym_utils.py index 007bf7a5..2b6157e6 100644 --- a/embodichain/lab/gym/utils/gym_utils.py +++ b/embodichain/lab/gym/utils/gym_utils.py @@ -42,6 +42,33 @@ "embodichain.lab.gym.envs.managers.rewards", ] +# Extra manager modules registered by third-party packages via init hooks +_EXTRA_MANAGER_MODULES: list[str] = [] + + +def register_manager_modules(modules: list[str]) -> None: + """Register additional manager modules for functor resolution. + + These modules are searched when resolving functor function names from + config strings, in addition to the built-in ``DEFAULT_MANAGER_MODULES``. + Call this from an ``embodichain.init`` entry point hook. + + Args: + modules: List of fully-qualified module path strings. + """ + for m in modules: + if m not in _EXTRA_MANAGER_MODULES: + _EXTRA_MANAGER_MODULES.append(m) + + +def get_manager_modules() -> list[str]: + """Get all registered manager modules (built-in + extensions). + + Returns: + Combined list of default and extra manager module paths. + """ + return DEFAULT_MANAGER_MODULES + _EXTRA_MANAGER_MODULES + def get_dtype_bounds(dtype: np.dtype): """Gets the min and max values of a given numpy type""" @@ -483,8 +510,8 @@ class ComponentCfg: env_cfg.sim_steps_per_control = config["env"].get("sim_steps_per_control", 4) env_cfg.extensions = deepcopy(config.get("env", {}).get("extensions", {})) - # Initialize manager_modules with defaults - default_manager_modules = DEFAULT_MANAGER_MODULES.copy() + # Initialize manager_modules with defaults + registered extensions + default_manager_modules = get_manager_modules().copy() # Extend with user-provided modules, skipping duplicates if manager_modules is not None: @@ -872,7 +899,7 @@ def build_env_cfg_from_args( gym_config = merge_args_with_gym_config(args, gym_config) cfg: EmbodiedEnvCfg = config_to_cfg( - gym_config, manager_modules=DEFAULT_MANAGER_MODULES + gym_config, manager_modules=get_manager_modules() ) cfg.filter_visual_rand = args.filter_visual_rand cfg.filter_dataset_saving = args.filter_dataset_saving diff --git a/embodichain/lab/gym/utils/registration.py b/embodichain/lab/gym/utils/registration.py index 9a5ae2af..a4bcf4fc 100644 --- a/embodichain/lab/gym/utils/registration.py +++ b/embodichain/lab/gym/utils/registration.py @@ -16,7 +16,10 @@ from __future__ import annotations +import importlib +import importlib.metadata import json +import logging import sys import torch @@ -33,6 +36,8 @@ if TYPE_CHECKING: from embodichain.lab.gym.envs import BaseEnv +_logger = logging.getLogger(__name__) + class EnvSpec: def __init__( @@ -200,3 +205,67 @@ def register_env_function(cls, uid, override=False, max_episode_steps=None, **kw kwargs=deepcopy(kwargs), ) return cls + + +def discover_task_packages() -> list[str]: + """Import all registered task packages via ``embodichain.tasks`` entry_points. + + Each task package's ``__init__.py`` recursively imports its sub-packages, + which triggers ``@register_env`` → ``gym.register()``. After this call, + all tasks from all installed packages are available in gymnasium's global + registry. + + Returns: + List of entry point names that were successfully imported. + """ + imported: list[str] = [] + try: + eps = importlib.metadata.entry_points(group="embodichain.tasks") + except TypeError: + # Python < 3.12: entry_points() requires a keyword argument + eps = importlib.metadata.entry_points().get("embodichain.tasks", []) + + for ep in eps: + try: + importlib.import_module(ep.value) + imported.append(ep.name) + except Exception: + _logger.warning( + f"Failed to import task package '{ep.name}' ({ep.value})", + exc_info=True, + ) + return imported + + +def execute_init_hooks() -> list[str]: + """Execute all registered init hooks via ``embodichain.init`` entry_points. + + Hooks are called in entry_points declaration order. An exception from one + hook does not prevent others from executing. + + Each entry point value must be in the format ``"module.path:function_name"``. + The function must accept no arguments and return ``None``. + + Returns: + List of hook names that were executed successfully. + """ + executed: list[str] = [] + try: + eps = importlib.metadata.entry_points(group="embodichain.init") + except TypeError: + # Python < 3.12 + eps = importlib.metadata.entry_points().get("embodichain.init", []) + + for ep in eps: + try: + module_name, func_name = ep.value.split(":", 1) + module = importlib.import_module(module_name) + func = getattr(module, func_name) + func() + executed.append(ep.name) + except Exception: + _logger.warning( + f"Init hook '{ep.name}' ({ep.value}) failed", + exc_info=True, + ) + return executed diff --git a/embodichain/lab/scripts/run_env.py b/embodichain/lab/scripts/run_env.py index 50ff124f..d34abfcb 100644 --- a/embodichain/lab/scripts/run_env.py +++ b/embodichain/lab/scripts/run_env.py @@ -25,6 +25,10 @@ add_env_launcher_args_to_parser, build_env_cfg_from_args, ) +from embodichain.lab.gym.utils.registration import ( + discover_task_packages, + execute_init_hooks, +) from embodichain.utils.logger import log_warning, log_info, log_error @@ -187,6 +191,12 @@ def cli(): args = parser.parse_args() + # Step 1: Discover all task packages via entry_points + discover_task_packages() + + # Step 2: Execute init hooks (register managers, asset resolvers, etc.) + execute_init_hooks() + env_cfg, gym_config, action_config = build_env_cfg_from_args(args) env = gymnasium.make(id=gym_config["id"], cfg=env_cfg, **action_config) diff --git a/embodichain_tasks/VERSION b/embodichain_tasks/VERSION new file mode 100644 index 00000000..8acdd82b --- /dev/null +++ b/embodichain_tasks/VERSION @@ -0,0 +1 @@ +0.0.1 diff --git a/embodichain_tasks/configs/agents/rl/basic/cart_pole/gym_config.json b/embodichain_tasks/configs/agents/rl/basic/cart_pole/gym_config.json new file mode 100644 index 00000000..085d94a3 --- /dev/null +++ b/embodichain_tasks/configs/agents/rl/basic/cart_pole/gym_config.json @@ -0,0 +1,74 @@ +{ + "id": "CartPoleRL", + "max_episodes": 5, + "max_episode_steps": 500, + "env": { + "events": {}, + "observations": { + "robot_qpos": { + "func": "normalize_robot_joint_data", + "mode": "modify", + "name": "robot/qpos", + "params": { + "joint_ids": [0, 1] + } + } + }, + "rewards": { + "velocity_penalty": { + "func": "joint_velocity_penalty", + "mode": "add", + "weight": 0.005, + "params": { + "robot_uid": "Cart", + "part_name": "hand" + } + } + }, + "actions": { + "delta_qpos": { + "func": "DeltaQposTerm", + "params": { "scale": 0.1 } + } + }, + "extensions": {} + }, + "robot": { + "uid": "Cart", + "urdf_cfg": { + "components": [ + { + "component_type": "arm", + "urdf_path": "CartPole/cart_pole.urdf" + } + ] + }, + "init_pos": [0.0, 0.0, 0.5], + "init_rot": [0.0, 0.0, 0.0], + "init_qpos": [-0.2, 0.07], + "drive_pros": { + "stiffness": { + "slider_to_cart": 1e1, + "cart_to_pole":1e-2 + }, + "damping": { + "slider_to_cart": 1e0, + "cart_to_pole":1e-3 + }, + "max_effort": { + "slider_to_cart": 1e2, + "cart_to_pole":1e-1 + } + }, + "control_parts": { + "arm": ["slider_to_cart"], + "hand": ["cart_to_pole"] + } + }, + "sensor": [], + "light": {}, + "background": [], + "rigid_object": [], + "rigid_object_group": [], + "articulation": [] +} diff --git a/embodichain_tasks/configs/agents/rl/basic/cart_pole/gym_config.yaml b/embodichain_tasks/configs/agents/rl/basic/cart_pole/gym_config.yaml new file mode 100644 index 00000000..e5d50843 --- /dev/null +++ b/embodichain_tasks/configs/agents/rl/basic/cart_pole/gym_config.yaml @@ -0,0 +1,66 @@ +id: CartPoleRL +max_episodes: 5 +max_episode_steps: 500 +env: + events: {} + observations: + robot_qpos: + func: normalize_robot_joint_data + mode: modify + name: robot/qpos + params: + joint_ids: + - 0 + - 1 + rewards: + velocity_penalty: + func: joint_velocity_penalty + mode: add + weight: 0.005 + params: + robot_uid: Cart + part_name: hand + actions: + delta_qpos: + func: DeltaQposTerm + params: + scale: 0.1 + extensions: {} +robot: + uid: Cart + urdf_cfg: + components: + - component_type: arm + urdf_path: CartPole/cart_pole.urdf + init_pos: + - 0.0 + - 0.0 + - 0.5 + init_rot: + - 0.0 + - 0.0 + - 0.0 + init_qpos: + - -0.2 + - 0.07 + drive_pros: + stiffness: + slider_to_cart: 10.0 + cart_to_pole: 0.01 + damping: + slider_to_cart: 1.0 + cart_to_pole: 0.001 + max_effort: + slider_to_cart: 100.0 + cart_to_pole: 0.1 + control_parts: + arm: + - slider_to_cart + hand: + - cart_to_pole +sensor: [] +light: {} +background: [] +rigid_object: [] +rigid_object_group: [] +articulation: [] diff --git a/embodichain_tasks/configs/agents/rl/basic/cart_pole/train_config.json b/embodichain_tasks/configs/agents/rl/basic/cart_pole/train_config.json new file mode 100644 index 00000000..6da5f735 --- /dev/null +++ b/embodichain_tasks/configs/agents/rl/basic/cart_pole/train_config.json @@ -0,0 +1,93 @@ +{ + "trainer": { + "exp_name": "cart_pole_ppo", + "gym_config": "configs/agents/rl/basic/cart_pole/gym_config.json", + "seed": 42, + "device": "cuda:0", + "headless": true, + "gpu_id": 0, + "num_envs": 64, + "iterations": 1000, + "buffer_size": 1024, + "eval_freq": 200, + "save_freq": 200, + "use_wandb": false, + "wandb_project_name": "embodichain-cart_pole", + "events": { + "eval": { + "record_camera": { + "func": "record_camera_data_async", + "mode": "interval", + "interval_step": 1, + "params": { + "name": "main_cam", + "resolution": [ + 640, + 480 + ], + "eye": [ + -1.4, + 1.4, + 2.5 + ], + "target": [ + 0, + 0, + 0.7 + ], + "up": [ + 0, + 0, + 1 + ], + "intrinsics": [ + 600, + 600, + 320, + 240 + ], + "save_path": "./outputs/videos/eval" + } + } + } + }, + "renderer": "fast-rt" + }, + "policy": { + "name": "actor_critic", + "actor": { + "type": "mlp", + "network_cfg": { + "hidden_sizes": [ + 256, + 256 + ], + "activation": "relu" + } + }, + "critic": { + "type": "mlp", + "network_cfg": { + "hidden_sizes": [ + 256, + 256 + ], + "activation": "relu" + } + } + }, + "algorithm": { + "name": "ppo", + "cfg": { + "learning_rate": 0.0001, + "n_epochs": 10, + "batch_size": 8192, + "gamma": 0.99, + "gae_lambda": 0.95, + "clip_coef": 0.2, + "ent_coef": 0.01, + "vf_coef": 0.5, + "max_grad_norm": 0.5 + } + } +} \ No newline at end of file diff --git a/embodichain_tasks/configs/agents/rl/basic/cart_pole/train_config.yaml b/embodichain_tasks/configs/agents/rl/basic/cart_pole/train_config.yaml new file mode 100644 index 00000000..b578bf96 --- /dev/null +++ b/embodichain_tasks/configs/agents/rl/basic/cart_pole/train_config.yaml @@ -0,0 +1,71 @@ +trainer: + exp_name: cart_pole_ppo + gym_config: configs/agents/rl/basic/cart_pole/gym_config.yaml + seed: 42 + device: cuda:0 + headless: true + num_envs: 64 + iterations: 1000 + buffer_size: 1024 + eval_freq: 200 + save_freq: 200 + use_wandb: false + wandb_project_name: embodichain-cart_pole + events: + eval: + record_camera: + func: record_camera_data_async + mode: interval + interval_step: 1 + params: + name: main_cam + resolution: + - 640 + - 480 + eye: + - -1.4 + - 1.4 + - 2.5 + target: + - 0 + - 0 + - 0.7 + up: + - 0 + - 0 + - 1 + intrinsics: + - 600 + - 600 + - 320 + - 240 + save_path: ./outputs/videos/eval + renderer: fast-rt +policy: + name: actor_critic + actor: + type: mlp + network_cfg: + hidden_sizes: + - 256 + - 256 + activation: relu + critic: + type: mlp + network_cfg: + hidden_sizes: + - 256 + - 256 + activation: relu +algorithm: + name: ppo + cfg: + learning_rate: 0.0001 + n_epochs: 10 + batch_size: 8192 + gamma: 0.99 + gae_lambda: 0.95 + clip_coef: 0.2 + ent_coef: 0.01 + vf_coef: 0.5 + max_grad_norm: 0.5 diff --git a/embodichain_tasks/configs/agents/rl/basic/cart_pole/train_config_grpo.json b/embodichain_tasks/configs/agents/rl/basic/cart_pole/train_config_grpo.json new file mode 100644 index 00000000..86ac34f2 --- /dev/null +++ b/embodichain_tasks/configs/agents/rl/basic/cart_pole/train_config_grpo.json @@ -0,0 +1,87 @@ +{ + "trainer": { + "exp_name": "cart_pole_grpo", + "gym_config": "configs/agents/rl/basic/cart_pole/gym_config.json", + "seed": 42, + "device": "cuda:0", + "headless": true, + "gpu_id": 0, + "num_envs": 64, + "iterations": 1000, + "buffer_size": 1024, + "eval_freq": 200, + "save_freq": 200, + "use_wandb": true, + "enable_eval": true, + "wandb_project_name": "embodichain-cart_pole", + "events": { + "eval": { + "record_camera": { + "func": "record_camera_data_async", + "mode": "interval", + "interval_step": 1, + "params": { + "name": "main_cam", + "resolution": [ + 640, + 480 + ], + "eye": [ + -1.4, + 1.4, + 2.5 + ], + "target": [ + 0, + 0, + 0.7 + ], + "up": [ + 0, + 0, + 1 + ], + "intrinsics": [ + 600, + 600, + 320, + 240 + ], + "save_path": "./outputs/videos/eval" + } + } + } + }, + "renderer": "hybrid" + }, + "policy": { + "name": "actor_only", + "actor": { + "type": "mlp", + "network_cfg": { + "hidden_sizes": [ + 256, + 256 + ], + "activation": "relu" + } + } + }, + "algorithm": { + "name": "grpo", + "cfg": { + "learning_rate": 0.0001, + "n_epochs": 10, + "batch_size": 8192, + "gamma": 0.99, + "clip_coef": 0.2, + "ent_coef": 0.01, + "kl_coef": 0.0, + "group_size": 4, + "eps": 1e-08, + "reset_every_rollout": true, + "max_grad_norm": 0.5, + "truncate_at_first_done": true + } + } +} \ No newline at end of file diff --git a/embodichain_tasks/configs/agents/rl/basic/cart_pole/train_config_grpo.yaml b/embodichain_tasks/configs/agents/rl/basic/cart_pole/train_config_grpo.yaml new file mode 100644 index 00000000..9b1966e3 --- /dev/null +++ b/embodichain_tasks/configs/agents/rl/basic/cart_pole/train_config_grpo.yaml @@ -0,0 +1,68 @@ +trainer: + exp_name: cart_pole_grpo + gym_config: configs/agents/rl/basic/cart_pole/gym_config.yaml + seed: 42 + device: cuda:0 + headless: true + num_envs: 64 + iterations: 1000 + buffer_size: 1024 + eval_freq: 200 + save_freq: 200 + use_wandb: true + enable_eval: true + wandb_project_name: embodichain-cart_pole + events: + eval: + record_camera: + func: record_camera_data_async + mode: interval + interval_step: 1 + params: + name: main_cam + resolution: + - 640 + - 480 + eye: + - -1.4 + - 1.4 + - 2.5 + target: + - 0 + - 0 + - 0.7 + up: + - 0 + - 0 + - 1 + intrinsics: + - 600 + - 600 + - 320 + - 240 + save_path: ./outputs/videos/eval + renderer: hybrid +policy: + name: actor_only + actor: + type: mlp + network_cfg: + hidden_sizes: + - 256 + - 256 + activation: relu +algorithm: + name: grpo + cfg: + learning_rate: 0.0001 + n_epochs: 10 + batch_size: 8192 + gamma: 0.99 + clip_coef: 0.2 + ent_coef: 0.01 + kl_coef: 0.0 + group_size: 4 + eps: 1.0e-08 + reset_every_rollout: true + max_grad_norm: 0.5 + truncate_at_first_done: true diff --git a/embodichain_tasks/configs/agents/rl/push_cube/gym_config.json b/embodichain_tasks/configs/agents/rl/push_cube/gym_config.json new file mode 100644 index 00000000..1399faf0 --- /dev/null +++ b/embodichain_tasks/configs/agents/rl/push_cube/gym_config.json @@ -0,0 +1,185 @@ +{ + "id": "PushCubeRL", + "max_episodes": 5, + "max_episode_steps": 100, + "env": { + "events": { + "randomize_cube": { + "func": "randomize_rigid_object_pose", + "mode": "reset", + "params": { + "entity_cfg": { + "uid": "cube" + }, + "position_range": [ + [-0.2, -0.2, 0.0], + [0.2, 0.2, 0.0] + ], + "relative_position": true + } + }, + "randomize_goal": { + "func": "randomize_target_pose", + "mode": "reset", + "params": { + "position_range": [ + [-0.3, -0.3, 0.05], + [0.3, 0.3, 0.05] + ], + "relative_position": false, + "store_key": "goal_pose" + } + } + }, + "observations": { + "robot_qpos": { + "func": "normalize_robot_joint_data", + "mode": "modify", + "name": "robot/qpos", + "params": { + "joint_ids": [0, 1, 2, 3, 4, 5] + } + }, + "robot_ee_pos": { + "func": "get_robot_eef_pose", + "mode": "add", + "name": "robot/ee_pos", + "params": { + "part_name": "arm" + } + }, + "cube_pos": { + "func": "get_rigid_object_pose", + "mode": "add", + "name": "object/cube_pos", + "params": { + "entity_cfg": { + "uid": "cube" + } + } + }, + "goal_pos": { + "func": "target_position", + "mode": "add", + "name": "object/goal_pos", + "params": { + "target_pose_key": "goal_pose" + } + } + }, + "rewards": { + "reaching_reward": { + "func": "reaching_behind_object", + "mode": "add", + "weight": 0.03, + "params": { + "object_cfg": { + "uid": "cube" + }, + "target_pose_key": "goal_pose", + "behind_offset": 0.03, + "height_offset": 0.015, + "distance_scale": 8.0, + "part_name": "arm" + } + }, + "goal_distance_reward": { + "func": "distance_to_target", + "mode": "add", + "weight": 0.8, + "params": { + "source_entity_cfg": { + "uid": "cube" + }, + "target_pose_key": "goal_pose", + "exponential": true, + "sigma": 0.12, + "use_xy_only": true + } + }, + + "action_penalty": { + "func": "action_smoothness_penalty", + "mode": "add", + "weight": 0.01, + "params": {} + }, + "success_bonus": { + "func": "success_reward", + "mode": "add", + "weight": 10.0, + "params": {} + } + }, + "actions": { + "delta_qpos": { + "func": "DeltaQposTerm", + "params": { "scale": 0.1 } + } + }, + "extensions": { + "success_threshold": 0.1 + } + }, + "robot": { + "robot_type": "URRobot", + "uid": "Manipulator", + "urdf_cfg": { + "components": [ + { + "component_type": "hand", + "urdf_path": "DH_PGI_140_80/DH_PGI_140_80.urdf" + } + ] + }, + "init_pos": [0.0, 0.0, 0.0], + "init_rot": [0.0, 0.0, 0.0], + "init_qpos": [0.0, -1.57, 1.57, -1.57, -1.57, 0.0, 0.04, 0.04], + "drive_pros": { + "drive_type": "force", + "stiffness": 100000.0, + "damping": 1000.0, + "max_velocity": 2.0, + "max_effort": 500.0 + }, + "solver_cfg": { + "arm": { + "tcp": [ + [1.0, 0.0, 0.0, 0.0], + [0.0, 1.0, 0.0, 0.0], + [0.0, 0.0, 1.0, 0.16], + [0.0, 0.0, 0.0, 1.0] + ] + } + } + }, + "sensor": [], + "light": {}, + "background": [], + "rigid_object": [ + { + "uid": "cube", + "shape": { + "shape_type": "Cube", + "size": [0.1, 0.1, 0.1] + }, + "body_type": "dynamic", + "init_pos": [-0.6, -0.4, 0.05], + "attrs": { + "mass": 2.0, + "static_friction": 1.0, + "dynamic_friction": 0.8, + "linear_damping": 2.0, + "angular_damping": 2.0, + "contact_offset": 0.003, + "rest_offset": 0.001, + "restitution": 0.1, + "max_depenetration_velocity": 10.0, + "max_linear_velocity": 1.0, + "max_angular_velocity": 1.0 + } + } + ], + "rigid_object_group": [], + "articulation": [] +} diff --git a/embodichain_tasks/configs/agents/rl/push_cube/train_config.json b/embodichain_tasks/configs/agents/rl/push_cube/train_config.json new file mode 100644 index 00000000..11b0972d --- /dev/null +++ b/embodichain_tasks/configs/agents/rl/push_cube/train_config.json @@ -0,0 +1,76 @@ +{ + "trainer": { + "exp_name": "push_cube_ppo", + "gym_config": "configs/agents/rl/push_cube/gym_config.json", + "seed": 42, + "device": "cuda:0", + "headless": true, + "gpu_id": 0, + "num_envs": 64, + "iterations": 1000, + "buffer_size": 1024, + "enable_eval": true, + "num_eval_envs": 16, + "num_eval_episodes": 3, + "eval_freq": 100, + "save_freq": 100, + "use_wandb": true, + "wandb_project_name": "embodichain-push_cube", + "events": { + "eval": { + "record_camera": { + "func": "record_camera_data_async", + "mode": "interval", + "interval_step": 1, + "params": { + "name": "main_cam", + "resolution": [640, 480], + "eye": [-1.4, 1.4, 2.0], + "target": [0, 0, 0], + "up": [0, 0, 1], + "intrinsics": [600, 600, 320, 240], + "save_path": "./outputs/videos_ppo1/eval" + } + } + } + }, + "renderer": "hybrid" + }, + "policy": { + "name": "actor_critic", + "actor": { + "type": "mlp", + "network_cfg": { + "hidden_sizes": [ + 256, + 256 + ], + "activation": "relu" + } + }, + "critic": { + "type": "mlp", + "network_cfg": { + "hidden_sizes": [ + 256, + 256 + ], + "activation": "relu" + } + } + }, + "algorithm": { + "name": "ppo", + "cfg": { + "learning_rate": 0.0001, + "n_epochs": 10, + "batch_size": 8192, + "gamma": 0.99, + "gae_lambda": 0.95, + "clip_coef": 0.2, + "ent_coef": 0.01, + "vf_coef": 0.5, + "max_grad_norm": 0.5 + } + } +} \ No newline at end of file diff --git a/embodichain_tasks/configs/agents/rl/push_cube/train_config_grpo.json b/embodichain_tasks/configs/agents/rl/push_cube/train_config_grpo.json new file mode 100644 index 00000000..df5f6681 --- /dev/null +++ b/embodichain_tasks/configs/agents/rl/push_cube/train_config_grpo.json @@ -0,0 +1,65 @@ +{ + "trainer": { + "exp_name": "push_cube_grpo", + "gym_config": "configs/agents/rl/push_cube/gym_config.json", + "seed": 42, + "device": "cuda:0", + "headless": true, + "gpu_id": 0, + "num_envs": 64, + "iterations": 1000, + "buffer_size": 1024, + "enable_eval": true, + "num_eval_envs": 16, + "num_eval_episodes": 3, + "eval_freq": 200, + "save_freq": 200, + "use_wandb": false, + "wandb_project_name": "embodichain-push_cube", + "events": { + "eval": { + "record_camera": { + "func": "record_camera_data_async", + "mode": "interval", + "interval_step": 1, + "params": { + "name": "main_cam", + "resolution": [640, 480], + "eye": [-1.4, 1.4, 2.0], + "target": [0, 0, 0], + "up": [0, 0, 1], + "intrinsics": [600, 600, 320, 240], + "save_path": "./outputs/videos/eval" + } + } + } + } + }, + "policy": { + "name": "actor_only", + "actor": { + "type": "mlp", + "network_cfg": { + "hidden_sizes": [256, 256], + "activation": "relu" + } + } + }, + "algorithm": { + "name": "grpo", + "cfg": { + "learning_rate": 0.0001, + "n_epochs": 10, + "batch_size": 8192, + "gamma": 0.99, + "clip_coef": 0.2, + "ent_coef": 0.01, + "kl_coef": 0.0, + "group_size": 4, + "eps": 1e-8, + "reset_every_rollout": true, + "max_grad_norm": 0.5, + "truncate_at_first_done": true + } + } +} diff --git a/embodichain_tasks/configs/gym/action_bank/conf.json b/embodichain_tasks/configs/gym/action_bank/conf.json new file mode 100644 index 00000000..3c5e9850 --- /dev/null +++ b/embodichain_tasks/configs/gym/action_bank/conf.json @@ -0,0 +1,237 @@ +{ + "scope": { + "right_arm": { + "type": "DiGraph", + "dim": [ + 6 + ], + "init": { + "method": "given_qpos", + "kwargs": { + "given_qpos": [ + 0, + 0, + 0, + 0, + 0, + 0 + ] + }, + "init_node_name": "" + }, + "dtype": "float32" + }, + "left_arm": { + "type": "DiGraph", + "dim": [ + 6 + ], + "init": { + "method": "given_qpos", + "kwargs": { + "given_qpos": [ + 0, + 0, + 0, + 0, + 0, + 0 + ] + }, + "init_node_name": "" + }, + "dtype": "float32" + }, + "left_eef": { + "type": "DiGraph", + "dim": [ + 2 + ], + "init": { + "method": "given_qpos", + "kwargs": { + "given_qpos": [ + 0, + 0 + ] + }, + "init_node_name": "" + }, + "dtype": "float32" + }, + "right_eef": { + "type": "DiGraph", + "dim": [ + 2 + ], + "init": { + "method": "given_qpos", + "kwargs": { + "given_qpos": [ + 0, + 0 + ] + }, + "init_node_name": "" + }, + "dtype": "float32" + } + }, + "node": { + "right_arm": [ + { + "home_qpos": { + "name": "A", + "kwargs": {} + } + }, + { + "bottle_grasp": { + "name": "B", + "kwargs": {} + } + }, + { + "bottle_pre1_pose": { + "name": "C", + "kwargs": {} + } + }, + { + "bottle_pre2_pose": { + "name": "C", + "kwargs": {} + } + }, + { + "bottle_place_pose": { + "name": "D", + "kwargs": {} + } + } + ], + "left_arm": [ + { + "lhome_qpos": { + "name": "a", + "kwargs": {} + } + }, + { + "cup_monitor_pose": { + "name": "b", + "kwargs": {} + } + } + ], + "left_eef": [ + { + "open": { + "name": "aa", + "kwargs": {} + } + }, + { + "close": { + "name": "bb", + "kwargs": {} + } + } + ], + "right_eef": [ + { + "ropen": { + "name": "cc", + "kwargs": {} + } + }, + { + "rclose": { + "name": "dd", + "kwargs": {} + } + } + ] + }, + "edge": { + "right_arm": [ + { + "init_to_pre1": { + "src": "home_qpos", + "sink": "bottle_pre1_pose", + "duration": 1, + "kwargs": {} + } + }, + { + "grasp_to_move": { + "src": "bottle_pre1_pose", + "sink": "bottle_grasp", + "duration": 2, + "kwargs": {} + } + }, + { + "move_to_rotation": { + "src": "bottle_grasp", + "sink": "bottle_pre2_pose", + "duration": 3, + "kwargs": {} + } + }, + { + "rotation_back_to_move": { + "src": "bottle_pre2_pose", + "sink": "bottle_place_pose", + "duration": 4, + "kwargs": {} + } + } + ], + "left_arm": [ + { + "init_to_monitor": { + "src": "lhome_qpos", + "sink": "cup_monitor_pose", + "duration": 1, + "kwargs": {} + } + }, + { + "left_arm_go_back": { + "src": "cup_monitor_pose", + "sink": "lhome_qpos", + "duration": 2, + "kwargs": {} + } + } + ], + "left_eef": [ + { + "lopen": { + "src": "close", + "sink": "open", + "duration": 10, + "kwargs": {} + } + } + ], + "right_eef": [ + { + "ropen": { + "src": "rclose", + "sink": "ropen", + "duration": 10, + "kwargs": {} + } + } + ] + }, + "sync": { + "grasp_to_move": { + "depend_tasks": [ + "init_to_monitor" + ] + } + } +} \ No newline at end of file diff --git a/embodichain_tasks/configs/gym/agent/pour_water_agent/agent_config.json b/embodichain_tasks/configs/gym/agent/pour_water_agent/agent_config.json new file mode 100644 index 00000000..b54f90c1 --- /dev/null +++ b/embodichain_tasks/configs/gym/agent/pour_water_agent/agent_config.json @@ -0,0 +1,31 @@ +{ "TaskAgent": { + "prompt_name": "one_stage_prompt" + }, + "CodeAgent": { + "prompt_name": "one_stage_prompt_according_to_task_plan" + }, + "Agent": { + "prompt_kwargs": { + "task_prompt": { + "type": "text", + "name": "task_prompt.txt" + }, + "basic_background": { + "type": "text", + "name": "basic_background.txt" + }, + "atom_actions": { + "type": "text", + "name": "atom_actions.txt" + }, + "code_prompt": { + "type": "text", + "name": "code_prompt.txt" + }, + "code_example": { + "type": "text", + "name": "code_example.txt" + } + } + } +} \ No newline at end of file diff --git a/embodichain_tasks/configs/gym/agent/pour_water_agent/agent_config_dual.json b/embodichain_tasks/configs/gym/agent/pour_water_agent/agent_config_dual.json new file mode 100644 index 00000000..b243bb65 --- /dev/null +++ b/embodichain_tasks/configs/gym/agent/pour_water_agent/agent_config_dual.json @@ -0,0 +1,31 @@ +{ "TaskAgent": { + "prompt_name": "one_stage_prompt" + }, + "CodeAgent": { + "prompt_name": "one_stage_prompt_according_to_task_plan" + }, + "Agent": { + "prompt_kwargs": { + "task_prompt": { + "type": "text", + "name": "task_prompt_dual.txt" + }, + "basic_background": { + "type": "text", + "name": "basic_background.txt" + }, + "atom_actions": { + "type": "text", + "name": "atom_actions.txt" + }, + "code_prompt": { + "type": "text", + "name": "code_prompt.txt" + }, + "code_example": { + "type": "text", + "name": "code_example.txt" + } + } + } +} \ No newline at end of file diff --git a/embodichain_tasks/configs/gym/agent/pour_water_agent/fast_gym_config.json b/embodichain_tasks/configs/gym/agent/pour_water_agent/fast_gym_config.json new file mode 100644 index 00000000..cdf0a685 --- /dev/null +++ b/embodichain_tasks/configs/gym/agent/pour_water_agent/fast_gym_config.json @@ -0,0 +1,391 @@ +{ + "id": "PourWaterAgent-v3", + "max_episodes": 5, + "env": { + "events": { + "init_table_pose": { + "func": "randomize_rigid_object_pose", + "mode": "reset", + "params": { + "entity_cfg": {"uid": "table"}, + "position_range": [[0.0, 0.0, -0.04], [0.0, 0.0, 0.04]], + "relative_position": true + } + }, + "init_bottle_pose": { + "func": "randomize_rigid_object_pose", + "mode": "reset", + "params": { + "entity_cfg": {"uid": "bottle"}, + "position_range": [[-0.08, -0.12, 0.0], [0.08, 0.04, 0.0]], + "relative_position": true + } + }, + "init_cup_pose": { + "func": "randomize_rigid_object_pose", + "mode": "reset", + "params": { + "entity_cfg": {"uid": "cup"}, + "position_range": [[-0.08, -0.04, 0.0], [0.04, 0.06, 0.0]], + "relative_position": true + } + }, + "prepare_extra_attr": { + "func": "prepare_extra_attr", + "mode": "reset", + "params": { + "attrs": [ + { + "name": "object_lengths", + "mode": "callable", + "entity_uids": "all_objects", + "func_name": "compute_object_length", + "func_kwargs": { + "is_svd_frame": true, + "sample_points": 5000 + } + }, + { + "name": "grasp_pose_object", + "mode": "static", + "entity_cfg": { + "uid": "bottle" + }, + "value": [[ + [0.32243, 0.03245, 0.94604, 0.025], + [0.00706, -0.99947, 0.03188, -0.0 ], + [0.94657, -0.0036 , -0.32249, 0.0 ], + [0.0 , 0.0 , 0.0 , 1.0 ] + ]] + }, + { + "name": "grasp_pose_object", + "mode": "static", + "entity_cfg": { + "uid": "cup" + }, + "value": [[ + [ 0.32039, -0.03227, 0.94673, 0.0 ], + [ 0.00675, -0.99932, -0.03635, 0.0 ], + [ 0.94726, 0.01803, -0.31996, 0.0 ], + [ 0.0 , 0.0 , 0.0 , 1.0 ] + ]] + }, + { + "name": "left_arm_base_pose", + "mode": "callable", + "entity_cfg": { + "uid": "CobotMagic" + }, + "func_name": "get_link_pose", + "func_kwargs": { + "link_name": "left_arm_base", + "to_matrix": true + } + }, + { + "name": "right_arm_base_pose", + "mode": "callable", + "entity_cfg": { + "uid": "CobotMagic" + }, + "func_name": "get_link_pose", + "func_kwargs": { + "link_name": "right_arm_base", + "to_matrix": true + } + } + ] + } + }, + "register_info_to_env": { + "func": "register_info_to_env", + "mode": "reset", + "params": { + "registry": [ + { + "entity_cfg": { + "uid": "bottle" + }, + "pose_register_params": { + "compute_relative": false, + "compute_pose_object_to_arena": true, + "to_matrix": true + } + }, + { + "entity_cfg": { + "uid": "cup" + }, + "pose_register_params": { + "compute_relative": false, + "compute_pose_object_to_arena": true, + "to_matrix": true + } + }, + { + "entity_cfg": { + "uid": "CobotMagic", + "control_parts": ["left_arm"] + }, + "attrs": ["left_arm_base_pose"], + "pose_register_params": { + "compute_relative": "cup", + "compute_pose_object_to_arena": false, + "to_matrix": true + }, + "prefix": false + }, + { + "entity_cfg": { + "uid": "CobotMagic", + "control_parts": ["right_arm"] + }, + "attrs": ["right_arm_base_pose"], + "pose_register_params": { + "compute_relative": "bottle", + "compute_pose_object_to_arena": false, + "to_matrix": true + }, + "prefix": false + } + ], + "registration": "affordance_datas", + "sim_update": true + } + }, + "random_table_material": { + "func": "randomize_visual_material", + "mode": "reset", + "interval_step": 2, + "params": { + "entity_cfg": {"uid": "table"}, + "random_texture_prob": 1.0, + "texture_path": "DexsimMaterials/WoodTable" + } + }, + "random_bottle_material": { + "func": "randomize_visual_material", + "mode": "reset", + "interval_step": 2, + "params": { + "entity_cfg": {"uid": "bottle"}, + "random_texture_prob": 1.0, + "texture_path": "DexsimMaterials/Bottle" + } + }, + "random_cup_material": { + "func": "randomize_visual_material", + "mode": "reset", + "interval_step": 2, + "params": { + "entity_cfg": {"uid": "cup"}, + "random_texture_prob": 1.0, + "texture_path": "DexsimMaterials/Cup" + } + }, + "random_robot_init_eef_pose": { + "func": "randomize_robot_eef_pose", + "mode": "reset", + "params": { + "entity_cfg": {"uid": "CobotMagic", "control_parts": ["left_arm", "right_arm"]}, + "position_range": [[-0.01, -0.01, -0.01], [0.01, 0.01, 0]] + } + }, + "record_camera": { + "func": "record_camera_data", + "mode": "interval", + "interval_step": 1, + "params": { + "name": "cam1", + "resolution": [320, 240], + "eye": [2, 0, 2], + "target": [0.5, 0, 1] + } + }, + "validation_cameras": { + "func": "validation_cameras", + "mode": "reset", + "params": { + "cameras": [ + { + "uid": "valid_cam_1", + "width": 1280, + "height": 960, + "enable_mask": false, + "intrinsics": [1400, 1400, 640, 480], + "extrinsics": { + "eye": [1.7, 0, 2.3], + "target": [0.6, 0, 0.8] + } + }, + { + "uid": "valid_cam_2", + "width": 1280, + "height": 960, + "enable_mask": false, + "intrinsics": [1400, 1400, 640, 480], + "extrinsics": { + "eye": [2.0, 0, 1.8], + "target": [0.7, 0, 0.9] + } + }, + { + "uid": "valid_cam_3", + "width": 1280, + "height": 960, + "enable_mask": false, + "intrinsics": [1400, 1400, 640, 480], + "extrinsics": { + "eye": [2.0, 0, 1.3], + "target": [0.7, 0, 0.9] + } + } + ] + } + } + }, + "dataset": { + "lerobot": { + "func": "LeRobotRecorder", + "mode": "save", + "params": { + "robot_meta": { + "control_freq": 25 + }, + "instruction": { + "lang": "Pour water from the bottle into the mug." + } + } + } + }, + "control_parts": ["left_arm", "left_eef", "right_arm", "right_eef"], + "success_params": { + "strict": false + } + }, + "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": "StereoCamera", + "uid": "cam_high", + "width": 960, + "height": 540, + "enable_mask": false, + "enable_depth": false, + "left_to_right_pos": [0.059684025824163614, 0, 0], + "intrinsics": [453.851402686215, 453.8347628855552, 469.827725021235, 258.6656181845155], + "intrinsics_right": [453.4536601653505, 453.3306024582175, 499.13697412367776, 297.7176248477935], + "extrinsics": { + "eye": [0.35368482807598, 0.014695524383058989, 1.4517046071614774], + "target": [0.7186357573287919, -0.054534732904795505, 0.5232553674540066], + "up": [0.9306678549330372, -0.0005600064212467153, 0.3658647703553347] + } + }, + { + "sensor_type": "Camera", + "uid": "cam_right_wrist", + "width": 640, + "height": 480, + "enable_mask": false, + "intrinsics": [488.1665344238281, 488.1665344238281, 322.7323303222656, 213.17434692382812], + "extrinsics": { + "parent": "right_link6", + "pos": [-0.08, 0.0, 0.04], + "quat": [0.15304635, 0.69034543, -0.69034543, -0.15304635] + } + }, + { + "sensor_type": "Camera", + "uid": "cam_left_wrist", + "width": 640, + "height": 480, + "enable_mask": false, + "intrinsics": [488.1665344238281, 488.1665344238281, 322.7323303222656, 213.17434692382812], + "extrinsics": { + "parent": "left_link6", + "pos": [-0.08, 0.0, 0.04], + "quat": [0.15304635, 0.69034543, -0.69034543, -0.15304635] + } + } + ], + "light": { + "direct": [ + { + "uid": "light_1", + "light_type": "point", + "color": [1.0, 1.0, 1.0], + "intensity": 100.0, + "init_pos": [2, 0, 2], + "radius": 20.0 + } + ] + }, + "background": [ + { + "uid": "table", + "shape": { + "shape_type": "Mesh", + "fpath": "CircleTableSimple/circle_table_simple.ply", + "compute_uv": true + }, + "attrs" : { + "mass": 10.0, + "static_friction": 0.95, + "dynamic_friction": 0.9, + "restitution": 0.01 + }, + "body_scale": [1, 1, 1], + "body_type": "kinematic", + "init_pos": [0.725, 0.0, 0.825], + "init_rot": [0, 90, 0] + } + ], + "rigid_object": [ + { + "uid":"cup", + "shape": { + "shape_type": "Mesh", + "fpath": "PaperCup/paper_cup.ply", + "compute_uv": true + }, + "attrs" : { + "mass": 0.01, + "contact_offset": 0.003, + "rest_offset": 0.001, + "restitution": 0.01, + "max_depenetration_velocity": 1e1, + "min_position_iters": 32, + "min_velocity_iters":8 + }, + "init_pos": [0.75, 0.1, 0.9], + "body_scale":[0.75, 0.75, 1.0], + "max_convex_hull_num": 8 + }, + { + "uid":"bottle", + "shape": { + "shape_type": "Mesh", + "fpath": "ScannedBottle/kashijia_processed.ply", + "compute_uv": true + }, + "attrs" : { + "mass": 0.01, + "contact_offset": 0.003, + "rest_offset": 0.001, + "restitution": 0.01, + "max_depenetration_velocity": 1e1, + "min_position_iters": 32, + "min_velocity_iters":8 + }, + "init_pos": [0.75, -0.1, 0.932], + "body_scale":[1, 1, 1], + "max_convex_hull_num": 8 + } + ] +} \ No newline at end of file diff --git a/embodichain_tasks/configs/gym/agent/pour_water_agent/task_prompt.txt b/embodichain_tasks/configs/gym/agent/pour_water_agent/task_prompt.txt new file mode 100644 index 00000000..be7eb3b1 --- /dev/null +++ b/embodichain_tasks/configs/gym/agent/pour_water_agent/task_prompt.txt @@ -0,0 +1,5 @@ +Task: +Use a single robotic arm to pour water from the bottle into the cup. +Position the bottle at an offset of [0.05, −0.10, 0.125] relative to the cup’s position during pouring. +After completing the pour, return the bottle to the location [0.7, −0.1]. +The cup should remain stationary throughout the task. \ No newline at end of file diff --git a/embodichain_tasks/configs/gym/agent/pour_water_agent/task_prompt_dual.txt b/embodichain_tasks/configs/gym/agent/pour_water_agent/task_prompt_dual.txt new file mode 100644 index 00000000..6b2205e2 --- /dev/null +++ b/embodichain_tasks/configs/gym/agent/pour_water_agent/task_prompt_dual.txt @@ -0,0 +1,5 @@ +Task: +Use both arms to pour water from the bottle into the cup. +First, grasp the bottle with right arm and the cup with left arm simultaneously. Then lift the cup by 0.10 m, and then move it to [0.55, 0.05] to prepare for pouring. +Then hold the bottle at a relative offset of [0.05, −0.10, 0.125] with respect to the cup for starting pouring. +After pouring, place the bottle to [0.7, −0.1] and place the cup at [0.6, 0.1] simultaneously. \ No newline at end of file diff --git a/embodichain_tasks/configs/gym/agent/rearrangement_agent/agent_config.json b/embodichain_tasks/configs/gym/agent/rearrangement_agent/agent_config.json new file mode 100644 index 00000000..6dbe2f5e --- /dev/null +++ b/embodichain_tasks/configs/gym/agent/rearrangement_agent/agent_config.json @@ -0,0 +1,31 @@ +{ "TaskAgent": { + "prompt_name": "one_stage_prompt" + }, + "CodeAgent": { + "prompt_name": "one_stage_prompt_according_to_task_plan" + }, + "Agent": { + "prompt_kwargs": { + "task_prompt": { + "type": "text", + "name": "task_prompt.txt" + }, + "basic_background": { + "type": "text", + "name": "basic_background.txt" + }, + "atom_actions": { + "type": "text", + "name": "atom_actions.txt" + }, + "code_prompt": { + "type": "text", + "name": "code_prompt.txt" + }, + "code_example": { + "type": "text", + "name": "code_example.txt" + } + } + } +} diff --git a/embodichain_tasks/configs/gym/agent/rearrangement_agent/fast_gym_config.json b/embodichain_tasks/configs/gym/agent/rearrangement_agent/fast_gym_config.json new file mode 100644 index 00000000..2fc603d0 --- /dev/null +++ b/embodichain_tasks/configs/gym/agent/rearrangement_agent/fast_gym_config.json @@ -0,0 +1,383 @@ +{ + "id": "RearrangementAgent-v3", + "max_episodes": 10, + "env": { + "events": { + "init_table_pose": { + "func": "randomize_rigid_object_pose", + "mode": "reset", + "params": { + "entity_cfg": {"uid": "table"}, + "position_range": [[0.0, 0.0, -0.04], [0.0, 0.0, 0.04]], + "relative_position": true + } + }, + "init_fork_pose": { + "func": "randomize_rigid_object_pose", + "mode": "reset", + "params": { + "entity_cfg": {"uid": "fork"}, + "position_range": [[0.0, -0.05, 0.0], [0.1, 0.05, 0.0]], + "rotation_range": [[0, 0, -45], [0, 0, 45]], + "relative_position": true, + "relative_rotation": true + } + }, + "init_spoon_pose": { + "func": "randomize_rigid_object_pose", + "mode": "reset", + "params": { + "entity_cfg": {"uid": "spoon"}, + "position_range": [[0.0, -0.05, 0.0], [0.1, 0.05, 0.0]], + "rotation_range": [[0, 0, -45], [0, 0, 45]], + "relative_position": true, + "relative_rotation": true + } + }, + "prepare_extra_attr": { + "func": "prepare_extra_attr", + "mode": "reset", + "params": { + "attrs": [ + { + "name": "object_lengths", + "mode": "callable", + "entity_uids": "all_objects", + "func_name": "compute_object_length", + "func_kwargs": { + "is_svd_frame": true, + "sample_points": 5000 + } + }, + { + "name": "grasp_pose_object", + "mode": "static", + "entity_uids": ["fork", "spoon"], + "value": [[ + [1.0, 0.0, 0.0, 0.0], + [0.0, -1.0, 0.0, 0.0], + [0.0, 0.0, -1.0, 0.0], + [0.0, 0.0, 0.0, 1.0] + ]] + } + ] + } + }, + "register_info_to_env": { + "func": "register_info_to_env", + "mode": "reset", + "params": { + "registry": [ + { + "entity_cfg": { + "uid": "plate" + }, + "pose_register_params": { + "compute_relative": false, + "compute_pose_object_to_arena": true, + "to_matrix": true + } + }, + { + "entity_cfg": { + "uid": "fork" + }, + "pose_register_params": { + "compute_relative": false, + "compute_pose_object_to_arena": true, + "to_matrix": true + } + }, + { + "entity_cfg": { + "uid": "spoon" + }, + "pose_register_params": { + "compute_relative": false, + "compute_pose_object_to_arena": true, + "to_matrix": true + } + }, + { + "entity_cfg": { + "uid": "CobotMagic", + "control_parts": ["left_arm"] + }, + "attrs": ["left_arm_base_pose"], + "pose_register_params": { + "compute_relative": "spoon", + "compute_pose_object_to_arena": false, + "to_matrix": true + }, + "prefix": false + }, + { + "entity_cfg": { + "uid": "CobotMagic", + "control_parts": ["right_arm"] + }, + "attrs": ["right_arm_base_pose"], + "pose_register_params": { + "compute_relative": "fork", + "compute_pose_object_to_arena": false, + "to_matrix": true + }, + "prefix": false + } + ], + "registration": "affordance_datas", + "sim_update": true + } + }, + "random_table_material": { + "func": "randomize_visual_material", + "mode": "reset", + "interval_step": 2, + "params": { + "entity_cfg": {"uid": "table"}, + "random_texture_prob": 1.0, + "texture_path": "DexsimMaterials/WoodTable" + } + }, + "random_plate_material": { + "func": "randomize_visual_material", + "mode": "reset", + "interval_step": 2, + "params": { + "entity_cfg": {"uid": "plate"}, + "random_texture_prob": 1.0, + "texture_path": "DexsimMaterials/Plate" + } + }, + "random_fork_material": { + "func": "randomize_visual_material", + "mode": "reset", + "interval_step": 2, + "params": { + "entity_cfg": {"uid": "fork"}, + "random_texture_prob": 1.0, + "texture_path": "DexsimMaterials/Spoon" + } + }, + "random_spoon_material": { + "func": "randomize_visual_material", + "mode": "reset", + "interval_step": 2, + "params": { + "entity_cfg": {"uid": "spoon"}, + "random_texture_prob": 1.0, + "texture_path": "DexsimMaterials/Spoon" + } + }, + "random_robot_init_eef_pose": { + "func": "randomize_robot_eef_pose", + "mode": "reset", + "params": { + "entity_cfg": {"uid": "CobotMagic", "control_parts": ["left_arm", "right_arm"]}, + "position_range": [[-0.01, -0.01, -0.01], [0.01, 0.01, 0]] + } + }, + "record_camera": { + "func": "record_camera_data", + "mode": "interval", + "interval_step": 1, + "params": { + "name": "cam1", + "resolution": [320, 240], + "eye": [2, 0, 2], + "target": [0.5, 0, 1] + } + }, + "validation_cameras": { + "func": "validation_cameras", + "mode": "trigger", + "params": { + "cameras": [ + { + "uid": "valid_cam_1", + "width": 1280, + "height": 960, + "enable_mask": false, + "intrinsics": [1400, 1400, 640, 480], + "extrinsics": { + "eye": [1.7, 0, 2.3], + "target": [0.6, 0, 0.8] + } + }, + { + "uid": "valid_cam_2", + "width": 1280, + "height": 960, + "enable_mask": false, + "intrinsics": [1400, 1400, 640, 480], + "extrinsics": { + "eye": [2.0, 0, 1.8], + "target": [0.7, 0, 0.9] + } + }, + { + "uid": "valid_cam_3", + "width": 1280, + "height": 960, + "enable_mask": false, + "intrinsics": [1400, 1400, 640, 480], + "extrinsics": { + "eye": [2.0, 0, 1.3], + "target": [0.7, 0, 0.9] + } + } + ] + } + } + }, + "dataset": { + "lerobot": { + "func": "LeRobotRecorder", + "mode": "save", + "params": { + "robot_meta": { + "control_freq": 25 + }, + "instruction": { + "lang": "Place the spoon and fork neatly into the plate on the table." + } + } + } + }, + "control_parts": ["left_arm", "left_eef", "right_arm", "right_eef"], + "success_params": { + "strict": false + } + }, + "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": "StereoCamera", + "uid": "cam_high", + "width": 960, + "height": 540, + "enable_mask": false, + "enable_depth": false, + "left_to_right_pos": [0.059684025824163614, 0, 0], + "intrinsics": [453.851402686215, 453.8347628855552, 469.827725021235, 258.6656181845155], + "intrinsics_right": [453.4536601653505, 453.3306024582175, 499.13697412367776, 297.7176248477935], + "extrinsics": { + "eye": [0.35368482807598, 0.014695524383058989, 1.4517046071614774], + "target": [0.7186357573287919, -0.054534732904795505, 0.5232553674540066], + "up": [0.9306678549330372, -0.0005600064212467153, 0.3658647703553347] + } + }, + { + "sensor_type": "Camera", + "uid": "cam_right_wrist", + "width": 640, + "height": 480, + "enable_mask": false, + "intrinsics": [488.1665344238281, 488.1665344238281, 322.7323303222656, 213.17434692382812], + "extrinsics": { + "parent": "right_link6", + "pos": [-0.08, 0.0, 0.04], + "quat": [0.15304635, 0.69034543, -0.69034543, -0.15304635] + } + }, + { + "sensor_type": "Camera", + "uid": "cam_left_wrist", + "width": 640, + "height": 480, + "enable_mask": false, + "intrinsics": [488.1665344238281, 488.1665344238281, 322.7323303222656, 213.17434692382812], + "extrinsics": { + "parent": "left_link6", + "pos": [-0.08, 0.0, 0.04], + "quat": [0.15304635, 0.69034543, -0.69034543, -0.15304635] + } + } + ], + "light": { + "direct": [ + { + "uid": "light_1", + "light_type": "point", + "color": [1.0, 1.0, 1.0], + "intensity": 100.0, + "init_pos": [2, 0, 2], + "radius": 20.0 + } + ] + }, + "background": [ + { + "uid": "table", + "shape": { + "shape_type": "Mesh", + "fpath": "CircleTableSimple/circle_table_simple.ply", + "compute_uv": true + }, + "attrs" : { + "mass": 10.0, + "static_friction": 2.0, + "dynamic_friction": 1.5, + "restitution": 0.1 + }, + "body_scale": [1, 1, 1], + "body_type": "kinematic", + "init_pos": [0.725, 0.0, 0.691], + "init_rot": [0, 90, 0] + } + ], + "rigid_object": [ + { + "uid": "plate", + "shape": { + "shape_type": "Mesh", + "fpath": "TableWare/tableware/plate/3_center.ply", + "compute_uv": true + }, + "body_scale": [0.001, 0.001, 0.001], + "init_pos": [0.5, 0.0, 1.0], + "init_rot": [180, 0, 0] + }, + { + "uid": "fork", + "shape": { + "shape_type": "Mesh", + "fpath": "TableWare/tableware/fork/standard_fork_scale.ply", + "compute_uv": true + }, + "body_scale": [1.0, 1.0, 1.0], + "init_pos": [0.5, 0.21, 1.0], + "attrs" : { + "mass": 0.01, + "static_friction": 0.95, + "dynamic_friction": 0.9, + "restitution": 0.05, + "min_position_iters": 16, + "min_velocity_iters": 4 + } + }, + { + "uid": "spoon", + "shape": { + "shape_type": "Mesh", + "fpath": "TableWare/tableware/spoon/standard_spoon_a_rescale.ply", + "compute_uv": true + }, + "body_scale": [1.0, 1.0, 1.0], + "init_pos": [0.5, -0.21, 1.0], + "attrs" : { + "mass": 0.01, + "static_friction": 0.95, + "dynamic_friction": 0.9, + "restitution": 0.05, + "min_position_iters": 16, + "min_velocity_iters": 4 + } + } + ] +} diff --git a/embodichain_tasks/configs/gym/agent/rearrangement_agent/task_prompt.txt b/embodichain_tasks/configs/gym/agent/rearrangement_agent/task_prompt.txt new file mode 100644 index 00000000..280bd227 --- /dev/null +++ b/embodichain_tasks/configs/gym/agent/rearrangement_agent/task_prompt.txt @@ -0,0 +1,8 @@ +Task: +Use both arms to rearrange a fork and a spoon on opposite sides of a plate. Perform the following steps **simultaneously**. + +1. Grasp the fork with the left arm and the spoon with the right arm. + +2. Reorient both end-effectors to a downward-facing pose. + +3. Place the fork at y = +0.16 and the spoon at y = −0.16 relative to the plate’s center. \ No newline at end of file diff --git a/embodichain_tasks/configs/gym/blocks_ranking_rgb/cobot_magic_3cam.json b/embodichain_tasks/configs/gym/blocks_ranking_rgb/cobot_magic_3cam.json new file mode 100644 index 00000000..5331ca4b --- /dev/null +++ b/embodichain_tasks/configs/gym/blocks_ranking_rgb/cobot_magic_3cam.json @@ -0,0 +1,283 @@ +{ + "id": "BlocksRankingRGB-v1", + "max_episodes": 5, + "env": { + "events": { + "random_light": { + "func": "randomize_light", + "mode": "interval", + "interval_step": 10, + "params": { + "entity_cfg": {"uid": "light_1"}, + "position_range": [[-0.5, -0.5, 2], [0.5, 0.5, 2]], + "color_range": [[0.6, 0.6, 0.6], [1, 1, 1]], + "intensity_range": [50.0, 100.0] + } + }, + "init_block_1_pose": { + "func": "randomize_rigid_object_pose", + "mode": "reset", + "params": { + "entity_cfg": {"uid": "block_1"}, + "position_range": [[0.60, -0.18, 0.9], [0.65, -0.12, 0.9]], + "rotation_range": [[0, 0, 0], [0, 0, 0]], + "relative_position": false + } + }, + "init_block_2_pose": { + "func": "randomize_rigid_object_pose", + "mode": "reset", + "params": { + "entity_cfg": {"uid": "block_2"}, + "position_range": [[0.65, -0.03, 0.9], [0.70, 0.03, 0.9]], + "rotation_range": [[0, 0, 0], [0, 0, 0]], + "relative_position": false + } + }, + "init_block_3_pose": { + "func": "randomize_rigid_object_pose", + "mode": "reset", + "params": { + "entity_cfg": {"uid": "block_3"}, + "position_range": [[0.70, 0.12, 0.9], [0.75, 0.18, 0.9]], + "rotation_range": [[0, 0, 0], [0, 0, 0]], + "relative_position": false + } + }, + "init_block_sizes": { + "func": "randomize_rigid_objects_scale", + "mode": "reset", + "params": { + "entity_cfgs": [ + {"uid": "block_1"}, + {"uid": "block_2"}, + {"uid": "block_3"} + ], + "scale_factor_range": [[0.75], [1.25]], + "same_scale_all_axes": true, + "shared_sample": true + } + }, + "set_block_1_color": { + "func": "set_rigid_object_visual_material", + "mode": "startup", + "params": { + "entity_cfg": {"uid": "block_1"}, + "mat_cfg": { + "uid": "red_block_mat", + "base_color": [1.0, 0.0, 0.0, 1.0], + "metallic": 0.0, + "roughness": 0.5 + } + } + }, + "set_block_2_color": { + "func": "set_rigid_object_visual_material", + "mode": "startup", + "params": { + "entity_cfg": {"uid": "block_2"}, + "mat_cfg": { + "uid": "green_block_mat", + "base_color": [0.0, 1.0, 0.0, 1.0], + "metallic": 0.0, + "roughness": 0.5 + } + } + }, + "set_block_3_color": { + "func": "set_rigid_object_visual_material", + "mode": "startup", + "params": { + "entity_cfg": {"uid": "block_3"}, + "mat_cfg": { + "uid": "blue_block_mat", + "base_color": [0.0, 0.0, 1.0, 1.0], + "metallic": 0.0, + "roughness": 0.5 + } + } + } + }, + "observations": { + "norm_robot_eef_joint": { + "func": "normalize_robot_joint_data", + "mode": "modify", + "name": "robot/qpos", + "params": { + "joint_ids": [6, 13] + } + }, + "block_1_pose": { + "func": "get_rigid_object_pose", + "mode": "add", + "name": "block_1_pose", + "params": { + "entity_cfg": {"uid": "block_1"} + } + }, + "block_2_pose": { + "func": "get_rigid_object_pose", + "mode": "add", + "name": "block_2_pose", + "params": { + "entity_cfg": {"uid": "block_2"} + } + }, + "block_3_pose": { + "func": "get_rigid_object_pose", + "mode": "add", + "name": "block_3_pose", + "params": { + "entity_cfg": {"uid": "block_3"} + } + } + } + }, + "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": "StereoCamera", + "uid": "cam_high", + "width": 960, + "height": 540, + "enable_mask": true, + "enable_depth": true, + "left_to_right_pos": [0.059684025824163614, 0, 0], + "intrinsics": [453.851402686215, 453.8347628855552, 469.827725021235, 258.6656181845155], + "intrinsics_right": [453.4536601653505, 453.3306024582175, 499.13697412367776, 297.7176248477935], + "extrinsics": { + "eye": [0.35368482807598, 0.014695524383058989, 1.4517046071614774], + "target": [0.7186357573287919, -0.054534732904795505, 0.5232553674540066], + "up": [0.9306678549330372, -0.0005600064212467153, 0.3658647703553347] + } + }, + { + "sensor_type": "Camera", + "uid": "cam_right_wrist", + "width": 640, + "height": 480, + "enable_mask": true, + "intrinsics": [488.1665344238281, 488.1665344238281, 322.7323303222656, 213.17434692382812], + "extrinsics": { + "parent": "right_link6", + "pos": [-0.08, 0.0, 0.04], + "quat": [0.15304635, 0.69034543, -0.69034543, -0.15304635] + } + }, + { + "sensor_type": "Camera", + "uid": "cam_left_wrist", + "width": 640, + "height": 480, + "enable_mask": true, + "intrinsics": [488.1665344238281, 488.1665344238281, 322.7323303222656, 213.17434692382812], + "extrinsics": { + "parent": "left_link6", + "pos": [-0.08, 0.0, 0.04], + "quat": [0.15304635, 0.69034543, -0.69034543, -0.15304635] + } + } + ], + "light": { + "direct": [ + { + "uid": "light_1", + "light_type": "point", + "color": [1.0, 1.0, 1.0], + "intensity": 50.0, + "init_pos": [2, 0, 2], + "radius": 10.0 + } + ] + }, + "background": [ + { + "uid": "table", + "shape": { + "shape_type": "Mesh", + "fpath": "CircleTableSimple/circle_table_simple.ply" + }, + "attrs" : { + "mass": 10.0, + "static_friction": 0.95, + "dynamic_friction": 0.9, + "restitution": 0.01 + }, + "body_scale": [1, 1, 1], + "body_type": "kinematic", + "init_pos": [0.725, 0.0, 0.825], + "init_rot": [0, 90, 0] + } + ], + "rigid_object": [ + { + "uid":"block_1", + "shape": { + "shape_type": "Cube", + "size": [0.04, 0.04, 0.04] + }, + "attrs" : { + "mass": 0.05, + "static_friction": 0.5, + "dynamic_friction": 0.5, + "restitution": 0.0, + "contact_offset": 0.003, + "rest_offset": 0.001, + "max_depenetration_velocity": 1e1, + "min_position_iters": 32, + "min_velocity_iters": 8 + }, + "init_pos": [0.725, -0.015, 0.86], + "body_scale":[1, 1, 1], + "max_convex_hull_num": 1 + }, + { + "uid":"block_2", + "shape": { + "shape_type": "Cube", + "size": [0.04, 0.04, 0.04] + }, + "attrs" : { + "mass": 0.05, + "static_friction": 0.5, + "dynamic_friction": 0.5, + "restitution": 0.0, + "contact_offset": 0.003, + "rest_offset": 0.001, + "max_depenetration_velocity": 1e1, + "min_position_iters": 32, + "min_velocity_iters": 8 + }, + "init_pos": [0.725, -0.015, 0.86], + "body_scale":[1, 1, 1], + "max_convex_hull_num": 1 + }, + { + "uid":"block_3", + "shape": { + "shape_type": "Cube", + "size": [0.04, 0.04, 0.04] + }, + "attrs" : { + "mass": 0.05, + "static_friction": 0.5, + "dynamic_friction": 0.5, + "restitution": 0.0, + "contact_offset": 0.003, + "rest_offset": 0.001, + "max_depenetration_velocity": 1e1, + "min_position_iters": 32, + "min_velocity_iters": 8 + }, + "init_pos": [0.725, -0.015, 0.86], + "body_scale":[1, 1, 1], + "max_convex_hull_num": 1 + } + ] +} + diff --git a/embodichain_tasks/configs/gym/blocks_ranking_size/cobot_magic_3cam.json b/embodichain_tasks/configs/gym/blocks_ranking_size/cobot_magic_3cam.json new file mode 100644 index 00000000..3f803066 --- /dev/null +++ b/embodichain_tasks/configs/gym/blocks_ranking_size/cobot_magic_3cam.json @@ -0,0 +1,257 @@ +{ + "id": "BlocksRankingSize-v1", + "max_episodes": 5, + "env": { + "events": { + "random_light": { + "func": "randomize_light", + "mode": "interval", + "interval_step": 10, + "params": { + "entity_cfg": {"uid": "light_1"}, + "position_range": [[-0.5, -0.5, 2], [0.5, 0.5, 2]], + "color_range": [[0.6, 0.6, 0.6], [1, 1, 1]], + "intensity_range": [50.0, 100.0] + } + }, + "init_block_1_pose": { + "func": "randomize_rigid_object_pose", + "mode": "reset", + "params": { + "entity_cfg": {"uid": "block_1"}, + "position_range": [[0.445, -0.08, 0.86], [1.005, 0.05, 0.86]], + "rotation_range": [[0, 0, -43.0], [0, 0, 43.0]], + "relative_position": false + } + }, + "init_block_2_pose": { + "func": "randomize_rigid_object_pose", + "mode": "reset", + "params": { + "entity_cfg": {"uid": "block_2"}, + "position_range": [[0.445, -0.08, 0.86], [1.005, 0.05, 0.86]], + "rotation_range": [[0, 0, -43.0], [0, 0, 43.0]], + "relative_position": false + } + }, + "init_block_3_pose": { + "func": "randomize_rigid_object_pose", + "mode": "reset", + "params": { + "entity_cfg": {"uid": "block_3"}, + "position_range": [[0.445, -0.08, 0.86], [1.005, 0.05, 0.86]], + "rotation_range": [[0, 0, -43.0], [0, 0, 43.0]], + "relative_position": false + } + }, + "init_block_1_size": { + "func": "randomize_rigid_object_scale", + "mode": "reset", + "params": { + "entity_cfg": {"uid": "block_1"}, + "scale_factor_range": [[0.95238095], [1.04761905]], + "same_scale_all_axes": true + } + }, + "init_block_2_size": { + "func": "randomize_rigid_object_scale", + "mode": "reset", + "params": { + "entity_cfg": {"uid": "block_2"}, + "scale_factor_range": [[0.94117647], [1.05882353]], + "same_scale_all_axes": true + } + }, + "init_block_3_size": { + "func": "randomize_rigid_object_scale", + "mode": "reset", + "params": { + "entity_cfg": {"uid": "block_3"}, + "scale_factor_range": [[0.92307692], [1.07692308]], + "same_scale_all_axes": true + } + } + }, + "observations": { + "norm_robot_eef_joint": { + "func": "normalize_robot_joint_data", + "mode": "modify", + "name": "robot/qpos", + "params": { + "joint_ids": [6, 13] + } + }, + "block_1_pose": { + "func": "get_rigid_object_pose", + "mode": "add", + "name": "block_1_pose", + "params": { + "entity_cfg": {"uid": "block_1"} + } + }, + "block_2_pose": { + "func": "get_rigid_object_pose", + "mode": "add", + "name": "block_2_pose", + "params": { + "entity_cfg": {"uid": "block_2"} + } + }, + "block_3_pose": { + "func": "get_rigid_object_pose", + "mode": "add", + "name": "block_3_pose", + "params": { + "entity_cfg": {"uid": "block_3"} + } + } + } + }, + "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": "StereoCamera", + "uid": "cam_high", + "width": 960, + "height": 540, + "enable_mask": true, + "enable_depth": true, + "left_to_right_pos": [0.059684025824163614, 0, 0], + "intrinsics": [453.851402686215, 453.8347628855552, 469.827725021235, 258.6656181845155], + "intrinsics_right": [453.4536601653505, 453.3306024582175, 499.13697412367776, 297.7176248477935], + "extrinsics": { + "eye": [0.35368482807598, 0.014695524383058989, 1.4517046071614774], + "target": [0.7186357573287919, -0.054534732904795505, 0.5232553674540066], + "up": [0.9306678549330372, -0.0005600064212467153, 0.3658647703553347] + } + }, + { + "sensor_type": "Camera", + "uid": "cam_right_wrist", + "width": 640, + "height": 480, + "enable_mask": true, + "intrinsics": [488.1665344238281, 488.1665344238281, 322.7323303222656, 213.17434692382812], + "extrinsics": { + "parent": "right_link6", + "pos": [-0.08, 0.0, 0.04], + "quat": [0.15304635, 0.69034543, -0.69034543, -0.15304635] + } + }, + { + "sensor_type": "Camera", + "uid": "cam_left_wrist", + "width": 640, + "height": 480, + "enable_mask": true, + "intrinsics": [488.1665344238281, 488.1665344238281, 322.7323303222656, 213.17434692382812], + "extrinsics": { + "parent": "left_link6", + "pos": [-0.08, 0.0, 0.04], + "quat": [0.15304635, 0.69034543, -0.69034543, -0.15304635] + } + } + ], + "light": { + "direct": [ + { + "uid": "light_1", + "light_type": "point", + "color": [1.0, 1.0, 1.0], + "intensity": 50.0, + "init_pos": [2, 0, 2], + "radius": 10.0 + } + ] + }, + "background": [ + { + "uid": "table", + "shape": { + "shape_type": "Mesh", + "fpath": "CircleTableSimple/circle_table_simple.ply" + }, + "attrs" : { + "mass": 10.0, + "static_friction": 0.95, + "dynamic_friction": 0.9, + "restitution": 0.01 + }, + "body_scale": [1, 1, 1], + "body_type": "kinematic", + "init_pos": [0.725, 0.0, 0.825], + "init_rot": [0, 90, 0] + } + ], + "rigid_object": [ + { + "uid":"block_1", + "shape": { + "shape_type": "Cube", + "size": [0.063, 0.063, 0.063] + }, + "attrs" : { + "mass": 0.05, + "static_friction": 0.5, + "dynamic_friction": 0.5, + "restitution": 0.0, + "contact_offset": 0.003, + "rest_offset": 0.001, + "max_depenetration_velocity": 1e1, + "min_position_iters": 32, + "min_velocity_iters": 8 + }, + "init_pos": [0.725, -0.015, 0.86], + "body_scale":[1, 1, 1], + "max_convex_hull_num": 1 + }, + { + "uid":"block_2", + "shape": { + "shape_type": "Cube", + "size": [0.051, 0.051, 0.051] + }, + "attrs" : { + "mass": 0.05, + "static_friction": 0.5, + "dynamic_friction": 0.5, + "restitution": 0.0, + "contact_offset": 0.003, + "rest_offset": 0.001, + "max_depenetration_velocity": 1e1, + "min_position_iters": 32, + "min_velocity_iters": 8 + }, + "init_pos": [0.725, -0.015, 0.86], + "body_scale":[1, 1, 1], + "max_convex_hull_num": 1 + }, + { + "uid":"block_3", + "shape": { + "shape_type": "Cube", + "size": [0.039, 0.039, 0.039] + }, + "attrs" : { + "mass": 0.05, + "static_friction": 0.5, + "dynamic_friction": 0.5, + "restitution": 0.0, + "contact_offset": 0.003, + "rest_offset": 0.001, + "max_depenetration_velocity": 1e1, + "min_position_iters": 32, + "min_velocity_iters": 8 + }, + "init_pos": [0.725, -0.015, 0.86], + "body_scale":[1, 1, 1], + "max_convex_hull_num": 1 + } + ] +} + diff --git a/embodichain_tasks/configs/gym/cobotmagic.json b/embodichain_tasks/configs/gym/cobotmagic.json new file mode 100644 index 00000000..a12c749a --- /dev/null +++ b/embodichain_tasks/configs/gym/cobotmagic.json @@ -0,0 +1,121 @@ +{ + "id": "EmbodiedEnv-v1", + "max_episodes": 10, + "env": { + "events": { + "random_light": { + "func": "randomize_light", + "mode": "interval", + "interval_step": 10, + "params": { + "entity_cfg": {"uid": "light_1"}, + "position_range": [[-0.5, -0.5, 2], [0.5, 0.5, 2]], + "color_range": [[0.6, 0.6, 0.6], [1, 1, 1]], + "intensity_range": [50.0, 100.0] + } + }, + "random_material": { + "func": "randomize_visual_material", + "mode": "interval", + "interval_step": 2, + "params": { + "entity_cfg": {"uid": "table"}, + "random_texture_prob": 0.5, + "texture_path": "CocoBackground/coco", + "base_color_range": [[0.2, 0.2, 0.2], [1.0, 1.0, 1.0]] + } + }, + "random_robot": { + "func": "randomize_visual_material", + "mode": "interval", + "interval_step": 5, + "params": { + "entity_cfg": {"uid": "CobotMagic", "link_names": [".*"]}, + "random_texture_prob": 0.5, + "texture_path": "CocoBackground/coco", + "base_color_range": [[0.2, 0.2, 0.2], [1.0, 1.0, 1.0]] + } + }, + "record_camera": { + "func": "record_camera_data", + "mode": "interval", + "interval_step": 1, + "params": { + "name": "cam1", + "resolution": [320, 240], + "eye": [2, 0, 2], + "target": [0.5, 0, 1] + } + }, + "replace_fork": { + "func": "replace_assets_from_group", + "mode": "reset", + "params": { + "entity_cfg": {"uid": "fork"}, + "folder_path": "TableWare/tableware/fork/" + } + } + } + }, + "sensor": [ + { + "uid": "camera_1", + "sensor_type": "Camera", + "width": 640, + "height": 480, + "enable_mask": true, + "enable_depth": true, + "extrinsics": { + "eye": [0.0, 0.0, 1.0], + "target": [0.0, 0.0, 0.0] + } + } + ], + "robot": { + "robot_type": "CobotMagic", + "init_pos": [0.0, 0.3, 1.2] + }, + "light": { + "direct": [ + { + "uid": "light_1", + "light_type": "point", + "color": [1.0, 1.0, 1.0], + "intensity": 50.0, + "init_pos": [0, 0, 2], + "radius": 10.0 + } + ] + }, + "background": [ + { + "uid": "table", + "shape": { + "shape_type": "Mesh", + "fpath": "ShopTableSimple/shop_table_simple.ply" + }, + "attrs" : { + "mass": 10.0 + }, + "body_scale": [2, 1.6, 1], + "body_type": "kinematic" + } + ], + "rigid_object": [ + { + "uid": "fork", + "shape": { + "shape_type": "Mesh", + "fpath": "TableWare/tableware/fork/standard_fork_scale.ply" + }, + "body_scale": [0.75, 0.75, 1.0], + "init_pos": [0.0, 0.0, 1.0] + } + ], + "articulation": [ + { + "fpath": "SlidingBoxDrawer/SlidingBoxDrawer.urdf", + "init_pos": [0.5, 0.0, 0.85] + } + ] +} diff --git a/embodichain_tasks/configs/gym/cobotmagic.yaml b/embodichain_tasks/configs/gym/cobotmagic.yaml new file mode 100644 index 00000000..f0ace0ce --- /dev/null +++ b/embodichain_tasks/configs/gym/cobotmagic.yaml @@ -0,0 +1,153 @@ +id: EmbodiedEnv-v1 +max_episodes: 10 +env: + events: + random_light: + func: randomize_light + mode: interval + interval_step: 10 + params: + entity_cfg: + uid: light_1 + position_range: + - - -0.5 + - -0.5 + - 2 + - - 0.5 + - 0.5 + - 2 + color_range: + - - 0.6 + - 0.6 + - 0.6 + - - 1 + - 1 + - 1 + intensity_range: + - 50.0 + - 100.0 + random_material: + func: randomize_visual_material + mode: interval + interval_step: 2 + params: + entity_cfg: + uid: table + random_texture_prob: 0.5 + texture_path: CocoBackground/coco + base_color_range: + - - 0.2 + - 0.2 + - 0.2 + - - 1.0 + - 1.0 + - 1.0 + random_robot: + func: randomize_visual_material + mode: interval + interval_step: 5 + params: + entity_cfg: + uid: CobotMagic + link_names: + - .* + random_texture_prob: 0.5 + texture_path: CocoBackground/coco + base_color_range: + - - 0.2 + - 0.2 + - 0.2 + - - 1.0 + - 1.0 + - 1.0 + record_camera: + func: record_camera_data + mode: interval + interval_step: 1 + params: + name: cam1 + resolution: + - 320 + - 240 + eye: + - 2 + - 0 + - 2 + target: + - 0.5 + - 0 + - 1 + replace_fork: + func: replace_assets_from_group + mode: reset + params: + entity_cfg: + uid: fork + folder_path: TableWare/tableware/fork/ +sensor: +- uid: camera_1 + sensor_type: Camera + width: 640 + height: 480 + enable_mask: true + enable_depth: true + extrinsics: + eye: + - 0.0 + - 0.0 + - 1.0 + target: + - 0.0 + - 0.0 + - 0.0 +robot: + robot_type: CobotMagic + init_pos: + - 0.0 + - 0.3 + - 1.2 +light: + direct: + - uid: light_1 + light_type: point + color: + - 1.0 + - 1.0 + - 1.0 + intensity: 50.0 + init_pos: + - 0 + - 0 + - 2 + radius: 10.0 +background: +- uid: table + shape: + shape_type: Mesh + fpath: ShopTableSimple/shop_table_simple.ply + attrs: + mass: 10.0 + body_scale: + - 2 + - 1.6 + - 1 + body_type: kinematic +rigid_object: +- uid: fork + shape: + shape_type: Mesh + fpath: TableWare/tableware/fork/standard_fork_scale.ply + body_scale: + - 0.75 + - 0.75 + - 1.0 + init_pos: + - 0.0 + - 0.0 + - 1.0 +articulation: +- fpath: SlidingBoxDrawer/SlidingBoxDrawer.urdf + init_pos: + - 0.5 + - 0.0 + - 0.85 diff --git a/embodichain_tasks/configs/gym/dexforce_w1.json b/embodichain_tasks/configs/gym/dexforce_w1.json new file mode 100644 index 00000000..42f0f073 --- /dev/null +++ b/embodichain_tasks/configs/gym/dexforce_w1.json @@ -0,0 +1,109 @@ +{ + "id": "EmbodiedEnv-v1", + "max_episodes": 10, + "env": { + "events": { + "random_light": { + "func": "randomize_light", + "mode": "interval", + "interval_step": 10, + "params": { + "entity_cfg": {"uid": "light_1"}, + "position_range": [[-0.5, -0.5, 2], [0.5, 0.5, 2]], + "color_range": [[0.6, 0.6, 0.6], [1, 1, 1]], + "intensity_range": [50.0, 100.0] + } + }, + "random_material": { + "func": "randomize_visual_material", + "mode": "interval", + "interval_step": 2, + "params": { + "entity_cfg": {"uid": "table"}, + "random_texture_prob": 0.5, + "texture_path": "CocoBackground/coco", + "base_color_range": [[0.2, 0.2, 0.2], [1.0, 1.0, 1.0]] + } + }, + "record_camera": { + "func": "record_camera_data", + "mode": "interval", + "interval_step": 1, + "params": { + "name": "cam1", + "resolution": [320, 240], + "eye": [2, 0, 2], + "target": [0.5, 0, 1] + } + }, + "replace_fork": { + "func": "replace_assets_from_group", + "mode": "reset", + "params": { + "entity_cfg": {"uid": "fork"}, + "folder_path": "TableWare/tableware/fork/" + } + } + } + }, + "sensor": [ + { + "sensor_type": "Camera", + "width": 640, + "height": 480, + "enable_mask": true, + "enable_depth": true, + "extrinsics": { + "eye": [0.0, 0.0, 1.0], + "target": [0.0, 0.0, 0.0] + } + } + ], + "robot": { + "robot_type": "DexforceW1", + "init_pos": [0.0, 1.0, 0] + }, + "light": { + "direct": [ + { + "uid": "light_1", + "light_type": "point", + "color": [1.0, 1.0, 1.0], + "intensity": 50.0, + "init_pos": [0, 0, 2], + "radius": 10.0 + } + ] + }, + "background": [ + { + "uid": "table", + "shape": { + "shape_type": "Mesh", + "fpath": "ShopTableSimple/shop_table_simple.ply" + }, + "attrs" : { + "mass": 10.0 + }, + "body_scale": [2, 1.6, 1], + "body_type": "kinematic" + } + ], + "rigid_object": [ + { + "uid": "fork", + "shape": { + "shape_type": "Mesh", + "fpath": "TableWare/tableware/fork/standard_fork_scale.ply" + }, + "body_scale": [0.75, 0.75, 1.0], + "init_pos": [0.0, 0.0, 1.0] + } + ], + "articulation": [ + { + "fpath": "SlidingBoxDrawer/SlidingBoxDrawer.urdf", + "init_pos": [0.5, 0.0, 0.85] + } + ] +} diff --git a/embodichain_tasks/configs/gym/dexforce_w1.yaml b/embodichain_tasks/configs/gym/dexforce_w1.yaml new file mode 100644 index 00000000..02f975b2 --- /dev/null +++ b/embodichain_tasks/configs/gym/dexforce_w1.yaml @@ -0,0 +1,134 @@ +id: EmbodiedEnv-v1 +max_episodes: 10 +env: + events: + random_light: + func: randomize_light + mode: interval + interval_step: 10 + params: + entity_cfg: + uid: light_1 + position_range: + - - -0.5 + - -0.5 + - 2 + - - 0.5 + - 0.5 + - 2 + color_range: + - - 0.6 + - 0.6 + - 0.6 + - - 1 + - 1 + - 1 + intensity_range: + - 50.0 + - 100.0 + random_material: + func: randomize_visual_material + mode: interval + interval_step: 2 + params: + entity_cfg: + uid: table + random_texture_prob: 0.5 + texture_path: CocoBackground/coco + base_color_range: + - - 0.2 + - 0.2 + - 0.2 + - - 1.0 + - 1.0 + - 1.0 + record_camera: + func: record_camera_data + mode: interval + interval_step: 1 + params: + name: cam1 + resolution: + - 320 + - 240 + eye: + - 2 + - 0 + - 2 + target: + - 0.5 + - 0 + - 1 + replace_fork: + func: replace_assets_from_group + mode: reset + params: + entity_cfg: + uid: fork + folder_path: TableWare/tableware/fork/ +sensor: +- sensor_type: Camera + width: 640 + height: 480 + enable_mask: true + enable_depth: true + extrinsics: + eye: + - 0.0 + - 0.0 + - 1.0 + target: + - 0.0 + - 0.0 + - 0.0 +robot: + robot_type: DexforceW1 + init_pos: + - 0.0 + - 1.0 + - 0 +light: + direct: + - uid: light_1 + light_type: point + color: + - 1.0 + - 1.0 + - 1.0 + intensity: 50.0 + init_pos: + - 0 + - 0 + - 2 + radius: 10.0 +background: +- uid: table + shape: + shape_type: Mesh + fpath: ShopTableSimple/shop_table_simple.ply + attrs: + mass: 10.0 + body_scale: + - 2 + - 1.6 + - 1 + body_type: kinematic +rigid_object: +- uid: fork + shape: + shape_type: Mesh + fpath: TableWare/tableware/fork/standard_fork_scale.ply + body_scale: + - 0.75 + - 0.75 + - 1.0 + init_pos: + - 0.0 + - 0.0 + - 1.0 +articulation: +- fpath: SlidingBoxDrawer/SlidingBoxDrawer.urdf + init_pos: + - 0.5 + - 0.0 + - 0.85 diff --git a/embodichain_tasks/configs/gym/match_object_container/cobot_magic_3cam.json b/embodichain_tasks/configs/gym/match_object_container/cobot_magic_3cam.json new file mode 100644 index 00000000..a127b47f --- /dev/null +++ b/embodichain_tasks/configs/gym/match_object_container/cobot_magic_3cam.json @@ -0,0 +1,343 @@ +{ + "id": "MatchObjectContainer-v1", + "max_episodes": 5, + "env": { + "events": { + "random_light": { + "func": "randomize_light", + "mode": "interval", + "interval_step": 10, + "params": { + "entity_cfg": {"uid": "light_1"}, + "position_range": [[-0.5, -0.5, 2], [0.5, 0.5, 2]], + "color_range": [[0.6, 0.6, 0.6], [1, 1, 1]], + "intensity_range": [50.0, 100.0] + } + }, + "init_block_cube_1_pose": { + "func": "randomize_rigid_object_pose", + "mode": "reset", + "params": { + "entity_cfg": {"uid": "block_cube_1"}, + "position_range": [[0.55, -0.10, 0.86], [0.58, -0.05, 0.86]], + "rotation_range": [[0, 0, -0.52], [0, 0, 0.52]], + "relative_position": false + } + }, + "init_block_sphere_1_pose": { + "func": "randomize_rigid_object_pose", + "mode": "reset", + "params": { + "entity_cfg": {"uid": "block_sphere_1"}, + "position_range": [[0.62, -0.10, 0.86], [0.65, -0.05, 0.86]], + "rotation_range": [[0, 0, -0.52], [0, 0, 0.52]], + "relative_position": false + } + }, + "init_block_cube_2_pose": { + "func": "randomize_rigid_object_pose", + "mode": "reset", + "params": { + "entity_cfg": {"uid": "block_cube_2"}, + "position_range": [[0.55, 0.05, 0.86], [0.58, 0.10, 0.86]], + "rotation_range": [[0, 0, -0.52], [0, 0, 0.52]], + "relative_position": false + } + }, + "init_block_sphere_2_pose": { + "func": "randomize_rigid_object_pose", + "mode": "reset", + "params": { + "entity_cfg": {"uid": "block_sphere_2"}, + "position_range": [[0.62, 0.05, 0.86], [0.65, 0.10, 0.86]], + "rotation_range": [[0, 0, -0.52], [0, 0, 0.52]], + "relative_position": false + } + } + }, + "observations": { + "norm_robot_eef_joint": { + "func": "normalize_robot_joint_data", + "mode": "modify", + "name": "robot/qpos", + "params": { + "joint_ids": [6, 13] + } + }, + "block_cube_1_pose": { + "func": "get_rigid_object_pose", + "mode": "add", + "name": "block_cube_1_pose", + "params": { + "entity_cfg": {"uid": "block_cube_1"} + } + }, + "block_sphere_1_pose": { + "func": "get_rigid_object_pose", + "mode": "add", + "name": "block_sphere_1_pose", + "params": { + "entity_cfg": {"uid": "block_sphere_1"} + } + }, + "container_cube_pose": { + "func": "get_rigid_object_pose", + "mode": "add", + "name": "container_cube_pose", + "params": { + "entity_cfg": {"uid": "container_cube"} + } + }, + "container_sphere_pose": { + "func": "get_rigid_object_pose", + "mode": "add", + "name": "container_sphere_pose", + "params": { + "entity_cfg": {"uid": "container_sphere"} + } + }, + "block_cube_2_pose": { + "func": "get_rigid_object_pose", + "mode": "add", + "name": "block_cube_2_pose", + "params": { + "entity_cfg": {"uid": "block_cube_2"} + } + }, + "block_sphere_2_pose": { + "func": "get_rigid_object_pose", + "mode": "add", + "name": "block_sphere_2_pose", + "params": { + "entity_cfg": {"uid": "block_sphere_2"} + } + } + } + }, + "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": "StereoCamera", + "uid": "cam_high", + "width": 960, + "height": 540, + "enable_mask": true, + "enable_depth": true, + "left_to_right_pos": [0.059684025824163614, 0, 0], + "intrinsics": [453.851402686215, 453.8347628855552, 469.827725021235, 258.6656181845155], + "intrinsics_right": [453.4536601653505, 453.3306024582175, 499.13697412367776, 297.7176248477935], + "extrinsics": { + "eye": [0.35368482807598, 0.014695524383058989, 1.4517046071614774], + "target": [0.7186357573287919, -0.054534732904795505, 0.5232553674540066], + "up": [0.9306678549330372, -0.0005600064212467153, 0.3658647703553347] + } + }, + { + "sensor_type": "Camera", + "uid": "cam_right_wrist", + "width": 640, + "height": 480, + "enable_mask": true, + "intrinsics": [488.1665344238281, 488.1665344238281, 322.7323303222656, 213.17434692382812], + "extrinsics": { + "parent": "right_link6", + "pos": [-0.08, 0.0, 0.04], + "quat": [0.15304635, 0.69034543, -0.69034543, -0.15304635] + } + }, + { + "sensor_type": "Camera", + "uid": "cam_left_wrist", + "width": 640, + "height": 480, + "enable_mask": true, + "intrinsics": [488.1665344238281, 488.1665344238281, 322.7323303222656, 213.17434692382812], + "extrinsics": { + "parent": "left_link6", + "pos": [-0.08, 0.0, 0.04], + "quat": [0.15304635, 0.69034543, -0.69034543, -0.15304635] + } + } + ], + "light": { + "direct": [ + { + "uid": "light_1", + "light_type": "point", + "color": [1.0, 1.0, 1.0], + "intensity": 50.0, + "init_pos": [2, 0, 2], + "radius": 10.0 + } + ] + }, + "background": [ + { + "uid": "table", + "shape": { + "shape_type": "Mesh", + "fpath": "CircleTableSimple/circle_table_simple.ply" + }, + "attrs" : { + "mass": 10.0, + "static_friction": 0.95, + "dynamic_friction": 0.9, + "restitution": 0.01 + }, + "body_scale": [1, 1, 1], + "body_type": "kinematic", + "init_pos": [0.725, 0.0, 0.825], + "init_rot": [0, 90, 0] + } + ], + "rigid_object": [ + { + "uid":"block_cube_1", + "shape": { + "shape_type": "Cube", + "size": [0.04, 0.04, 0.04] + }, + "attrs" : { + "mass": 0.05, + "static_friction": 0.5, + "dynamic_friction": 0.5, + "restitution": 0.0, + "contact_offset": 0.003, + "rest_offset": 0.001, + "max_depenetration_velocity": 1e1, + "min_position_iters": 32, + "min_velocity_iters": 8 + }, + "init_pos": [0.565, -0.075, 0.86], + "init_rot": [0, 0, 0], + "body_scale":[1, 1, 1], + "max_convex_hull_num": 1 + }, + { + "uid":"block_sphere_1", + "shape": { + "shape_type": "Sphere", + "radius": 0.025 + }, + "attrs" : { + "mass": 0.05, + "static_friction": 0.5, + "dynamic_friction": 0.5, + "restitution": 0.0, + "contact_offset": 0.003, + "rest_offset": 0.001, + "max_depenetration_velocity": 1e1, + "min_position_iters": 32, + "min_velocity_iters": 8 + }, + "init_pos": [0.635, -0.075, 0.86], + "init_rot": [0, 0, 0], + "body_scale":[1, 1, 1], + "max_convex_hull_num": 1 + }, + { + "uid":"container_cube", + "shape": { + "shape_type": "Mesh", + "fpath": "ContainerMetal/container_metal.obj" + }, + "body_type": "dynamic", + "attrs" : { + "mass": 0.5, + "static_friction": 1.0, + "dynamic_friction": 1.0, + "restitution": 0.0, + "contact_offset": 0.003, + "rest_offset": 0.001, + "max_depenetration_velocity": 2.0, + "linear_damping": 2.0, + "angular_damping": 2.0, + "max_linear_velocity": 5.0, + "max_angular_velocity": 10.0, + "min_position_iters": 32, + "min_velocity_iters": 8 + }, + "init_pos": [0.875, -0.25, 0.86], + "init_rot": [0, 0, 0], + "body_scale":[1, 1, 1], + "max_convex_hull_num": 8 + }, + { + "uid":"container_sphere", + "shape": { + "shape_type": "Mesh", + "fpath": "ContainerMetal/container_metal.obj" + }, + "body_type": "dynamic", + "attrs" : { + "mass": 0.5, + "static_friction": 1.0, + "dynamic_friction": 1.0, + "restitution": 0.0, + "contact_offset": 0.003, + "rest_offset": 0.001, + "max_depenetration_velocity": 2.0, + "linear_damping": 2.0, + "angular_damping": 2.0, + "max_linear_velocity": 5.0, + "max_angular_velocity": 10.0, + "min_position_iters": 32, + "min_velocity_iters": 8 + }, + "init_pos": [0.875, 0.25, 0.86], + "init_rot": [0, 0, 0], + "body_scale":[1, 1, 1], + "max_convex_hull_num": 8 + }, + { + "uid":"block_cube_2", + "shape": { + "shape_type": "Cube", + "size": [0.04, 0.04, 0.04] + }, + "attrs" : { + "mass": 0.05, + "static_friction": 0.5, + "dynamic_friction": 0.5, + "restitution": 0.0, + "contact_offset": 0.003, + "rest_offset": 0.001, + "max_depenetration_velocity": 1e1, + "min_position_iters": 32, + "min_velocity_iters": 8 + }, + "init_pos": [0.565, 0.075, 0.86], + "init_rot": [0, 0, 0], + "body_scale":[1, 1, 1], + "max_convex_hull_num": 1 + }, + { + "uid":"block_sphere_2", + "shape": { + "shape_type": "Sphere", + "radius": 0.025 + }, + "attrs" : { + "mass": 0.05, + "static_friction": 0.5, + "dynamic_friction": 0.5, + "restitution": 0.0, + "contact_offset": 0.003, + "rest_offset": 0.001, + "max_depenetration_velocity": 1e1, + "min_position_iters": 32, + "min_velocity_iters": 8 + }, + "init_pos": [0.635, 0.075, 0.86], + "init_rot": [0, 0, 0], + "body_scale":[1, 1, 1], + "max_convex_hull_num": 1 + } + ] +} + diff --git a/embodichain_tasks/configs/gym/place_object_drawer/cobot_magic_3cam.json b/embodichain_tasks/configs/gym/place_object_drawer/cobot_magic_3cam.json new file mode 100644 index 00000000..185e4617 --- /dev/null +++ b/embodichain_tasks/configs/gym/place_object_drawer/cobot_magic_3cam.json @@ -0,0 +1,165 @@ +{ + "id": "PlaceObjectDrawer-v1", + "max_episodes": 5, + "env": { + "events": { + "random_light": { + "func": "randomize_light", + "mode": "interval", + "interval_step": 10, + "params": { + "entity_cfg": {"uid": "light_1"}, + "position_range": [[-0.5, -0.5, 2], [0.5, 0.5, 2]], + "color_range": [[0.6, 0.6, 0.6], [1, 1, 1]], + "intensity_range": [50.0, 100.0] + } + }, + "init_object_pose": { + "func": "randomize_rigid_object_pose", + "mode": "reset", + "params": { + "entity_cfg": {"uid": "object"}, + "position_range": [[0.66, -0.15, 0.86], [0.70, -0.10, 0.86]], + "rotation_range": [[0, 0, -0.52], [0, 0, 0.52]], + "relative_position": false + } + } + }, + "observations": { + "norm_robot_eef_joint": { + "func": "normalize_robot_joint_data", + "mode": "modify", + "name": "robot/qpos", + "params": { + "joint_ids": [12, 13, 14, 15] + } + }, + "object_pose": { + "func": "get_rigid_object_pose", + "mode": "add", + "name": "object_pose", + "params": { + "entity_cfg": {"uid": "object"} + } + } + } + }, + "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": "StereoCamera", + "uid": "cam_high", + "width": 960, + "height": 540, + "enable_mask": true, + "enable_depth": true, + "left_to_right_pos": [0.059684025824163614, 0, 0], + "intrinsics": [453.851402686215, 453.8347628855552, 469.827725021235, 258.6656181845155], + "intrinsics_right": [453.4536601653505, 453.3306024582175, 499.13697412367776, 297.7176248477935], + "extrinsics": { + "eye": [0.35368482807598, 0.014695524383058989, 1.4517046071614774], + "target": [0.7186357573287919, -0.054534732904795505, 0.5232553674540066], + "up": [0.9306678549330372, -0.0005600064212467153, 0.3658647703553347] + } + }, + { + "sensor_type": "Camera", + "uid": "cam_right_wrist", + "width": 640, + "height": 480, + "enable_mask": true, + "intrinsics": [488.1665344238281, 488.1665344238281, 322.7323303222656, 213.17434692382812], + "extrinsics": { + "parent": "right_link6", + "pos": [-0.08, 0.0, 0.04], + "quat": [0.15304635, 0.69034543, -0.69034543, -0.15304635] + } + }, + { + "sensor_type": "Camera", + "uid": "cam_left_wrist", + "width": 640, + "height": 480, + "enable_mask": true, + "intrinsics": [488.1665344238281, 488.1665344238281, 322.7323303222656, 213.17434692382812], + "extrinsics": { + "parent": "left_link6", + "pos": [-0.08, 0.0, 0.04], + "quat": [0.15304635, 0.69034543, -0.69034543, -0.15304635] + } + } + ], + "light": { + "direct": [ + { + "uid": "light_1", + "light_type": "point", + "color": [1.0, 1.0, 1.0], + "intensity": 50.0, + "init_pos": [2, 0, 2], + "radius": 10.0 + } + ] + }, + "background": [ + { + "uid": "table", + "shape": { + "shape_type": "Mesh", + "fpath": "CircleTableSimple/circle_table_simple.ply" + }, + "attrs" : { + "mass": 10.0, + "static_friction": 0.95, + "dynamic_friction": 0.9, + "restitution": 0.01 + }, + "body_scale": [1, 1, 1], + "body_type": "kinematic", + "init_pos": [0.725, 0.0, 0.825], + "init_rot": [0, 90, 0] + } + ], + "rigid_object": [ + { + "uid":"object", + "shape": { + "shape_type": "Mesh", + "fpath": "ToyDuck/toy_duck.glb" + }, + "attrs" : { + "mass": 0.01, + "static_friction": 1.0, + "dynamic_friction": 1.0, + "restitution": 0.0, + "contact_offset": 0.003, + "rest_offset": 0.001, + "max_depenetration_velocity": 2.0, + "linear_damping": 2.0, + "angular_damping": 2.0, + "max_linear_velocity": 5.0, + "max_angular_velocity": 10.0, + "min_position_iters": 32, + "min_velocity_iters": 8 + }, + "init_pos": [0.725, -0.1, 0.86], + "init_rot": [0, 0, 0], + "body_scale":[0.2, 0.2, 0.2], + "max_convex_hull_num": 8 + } + ], + "articulation": [ + { + "uid": "drawer", + "fpath": "SimpleBoxDrawer/simple_box_drawer/simple_box_drawer.urdf", + "init_pos": [0.725, 0.16, 1.025], + "init_rot": [0, 180, 0] + } + ] +} + diff --git a/embodichain_tasks/configs/gym/pour_water/action_config.json b/embodichain_tasks/configs/gym/pour_water/action_config.json new file mode 100644 index 00000000..d49561a4 --- /dev/null +++ b/embodichain_tasks/configs/gym/pour_water/action_config.json @@ -0,0 +1,938 @@ +{ + "scope": { + "right_arm": { + "type": "DiGraph", + "dim": [ + 6 + ], + "init": { + "method": "current_qpos", + "init_node_name": "right_arm_init_qpos" + }, + "dtype": "float32" + }, + "left_arm": { + "type": "DiGraph", + "dim": [ + 6 + ], + "init": { + "method": "current_qpos", + "init_node_name": "left_arm_init_qpos" + }, + "dtype": "float32" + }, + "left_eef": { + "type": "DiGraph", + "dim": [ + 1 + ], + "init": { + "method": "given_qpos", + "kwargs": { + "given_qpos": [ + 1 + ] + }, + "init_node_name": "" + }, + "dtype": "float32" + }, + "right_eef": { + "type": "DiGraph", + "dim": [ + 1 + ], + "init": { + "method": "given_qpos", + "kwargs": { + "given_qpos": [ + 1 + ] + }, + "init_node_name": "" + }, + "dtype": "float32" + } + }, + "node": { + "right_arm": [ + { + "right_arm_init_qpos": { + "name": "generate_affordances_from_src", + "kwargs": { + "affordance_infos": [ + { + "src_key": "right_arm_init_qpos", + "dst_key": "right_arm_init_pose", + "valid_funcs_name_kwargs_proc": [ + { + "name": "no_validation", + "kwargs": {}, + "pass_processes": [ + { + "name": "get_fk_xpos", + "kwargs": { + "control_part": "right_arm", + "fk_func": "env.robot.compute_fk" + } + } + ] + } + ] + } + ] + + } + } + }, + { + "right_arm_aim_qpos": { + "name": "generate_right_arm_aim_qpos", + "kwargs": {} + } + }, + { + "bottle_grasp": { + "name": "generate_affordances_from_src", + "kwargs": { + "affordance_infos": [ + { + "src_key": "bottle_pose", + "dst_key": "bottle_grasp_pose", + "valid_funcs_name_kwargs_proc": [ + { + "name": "no_validation", + "kwargs": {}, + "pass_processes": [ + { + "name": "get_rotation_replaced_pose", + "kwargs": { + "rotation_value": "env.affordance_datas['right_arm_aim_qpos'][0]", + "rot_axis": "z", + "mode": "intrinsic" + } + }, + { + "name": "get_frame_changed_pose", + "kwargs": { + "frame_change_matrix": "env.affordance_datas['bottle_pose']", + "mode": "intrinsic", + "inverse": true + } + }, + { + "name": "get_frame_changed_pose", + "kwargs": { + "frame_change_matrix": "env.affordance_datas['bottle_grasp_pose']", + "mode": "intrinsic" + } + } + ] + } + ] + }, + { + "src_key": "bottle_grasp_pose", + "dst_key": "bottle_pre1_pose", + "valid_funcs_name_kwargs_proc": [ + { + "name": "no_validation", + "kwargs": {}, + "pass_processes": [ + { + "name": "get_offset_pose", + "kwargs": { + "offset_value": -0.05, + "direction": "z", + "mode": "intrinsic" + } + } + ] + } + ] + }, + { + "src_key": "bottle_pre1_pose", + "dst_key": "bottle_pre2_pose", + "valid_funcs_name_kwargs_proc": [ + { + "name": "no_validation", + "kwargs": {}, + "pass_processes": [ + { + "name": "get_offset_pose", + "kwargs": { + "offset_value": -0.05, + "direction": "z", + "mode": "intrinsic" + } + } + ] + } + ] + } + ] + } + } + }, + { + "bottle_pre1_qpos": { + "name": "generate_affordances_from_src", + "kwargs": { + "affordance_infos": [ + { + "src_key": "bottle_pre1_pose", + "dst_key": "bottle_pre1_qpos", + "valid_funcs_name_kwargs_proc": [ + { + "name": "get_ik_ret", + "kwargs": { + "ik_func": "env.robot.compute_ik", + "qpos_seed": "env.affordance_datas['right_arm_init_qpos']", + "control_part": "right_arm" + }, + "pass_processes": [ + { + "name": "get_ik_qpos", + "kwargs": { + "ik_func": "env.robot.compute_ik", + "qpos_seed": "env.affordance_datas['right_arm_init_qpos']", + "control_part": "right_arm" + } + } + ] + }, + { + "name": "is_qpos_flip", + "kwargs": { + "qpos_ref": "env.affordance_datas['right_arm_init_qpos']", + "qpos_ids": [3, 4], + "threshold": 3.455751918948773, + "mode": "delta", + "return_inverse": true + } + } + ] + } + ] + } + } + }, + { + "bottle_grasp_qpos": { + "name": "generate_affordances_from_src", + "kwargs": { + "affordance_infos": [ + { + "src_key": "bottle_grasp_pose", + "dst_key": "bottle_grasp_qpos", + "valid_funcs_name_kwargs_proc": [ + { + "name": "get_ik_ret", + "kwargs": { + "ik_func": "env.robot.compute_ik", + "qpos_seed": "env.affordance_datas['bottle_pre1_qpos']", + "control_part": "right_arm" + }, + "pass_processes": [ + { + "name": "get_ik_qpos", + "kwargs": { + "ik_func": "env.robot.compute_ik", + "qpos_seed": "env.affordance_datas['bottle_pre1_qpos']", + "control_part": "right_arm" + } + } + ] + }, + { + "name": "is_qpos_flip", + "kwargs": { + "qpos_ref": "env.affordance_datas['right_arm_init_qpos']", + "qpos_ids": [3, 4], + "threshold": 3.455751918948773, + "mode": "delta", + "return_inverse": true + } + } + ] + } + ] + + } + } + }, + { + "bottle_up_qpos": { + "name": "generate_affordances_from_src", + "kwargs": { + "affordance_infos": [ + { + "src_key": "bottle_grasp_qpos", + "dst_key": "bottle_up_qpos", + "valid_funcs_name_kwargs_proc": [ + { + "name": "no_validation", + "kwargs": {}, + "pass_processes": [ + { + "name": "get_offset_qpos", + "kwargs": { + "offset_value": -0.05, + "joint_list_offset": [1] + } + } + ] + }, + { + "name": "is_qpos_exceed", + "kwargs": { + "robot": "env.robot", + "control_part": "right_arm" + } + }, + { + "name": "is_qpos_flip", + "kwargs": { + "qpos_ref": "env.affordance_datas['right_arm_init_qpos']", + "qpos_ids": [3, 4], + "threshold": 3.455751918948773, + "mode": "delta", + "return_inverse": true + } + } + ] + } + ] + } + } + }, + { + "pour_water_start_qpos": { + "name": "generate_affordances_from_src", + "kwargs": { + "affordance_infos": [ + { + "src_key": "bottle_grasp_pose", + "dst_key": "pour_water_start_pose", + "valid_funcs_name_kwargs_proc": [ + { + "name":"no_validation", + "kwargs": {}, + "pass_processes": [ + { + "name": "get_replaced_pose", + "kwargs": { + "pose_replace_value": "env.affordance_datas['cup_pose'][:2,3]", + "axis_str_replace": "xy" + } + }, + { + "name": "get_offset_pose", + "kwargs": { + "offset_value": 0.05, + "direction": "x", + "mode": "extrinsic" + } + }, + { + "name": "get_offset_pose", + "kwargs": { + "offset_value": -0.10, + "direction": "y", + "mode": "extrinsic" + } + }, + { + "name": "get_offset_pose", + "kwargs": { + "offset_value": 0.125, + "direction": "z", + "mode": "extrinsic" + } + } + ] + } + ] + }, + { + "src_key": "pour_water_start_pose", + "dst_key": "pour_water_start_qpos", + "valid_funcs_name_kwargs_proc": [ + { + "name": "get_ik_ret", + "kwargs": { + "ik_func": "env.robot.compute_ik", + "qpos_seed": "env.affordance_datas['bottle_up_qpos']", + "control_part": "right_arm" + }, + "pass_processes": [ + { + "name": "get_ik_qpos", + "kwargs": { + "ik_func": "env.robot.compute_ik", + "qpos_seed": "env.affordance_datas['bottle_up_qpos']", + "control_part": "right_arm" + } + } + ] + }, + { + "name": "is_qpos_flip", + "kwargs": { + "qpos_ref": "env.affordance_datas['right_arm_init_qpos']", + "qpos_ids": [3, 4], + "threshold": 3.455751918948773, + "mode": "delta", + "return_inverse": true + } + } + ] + } + ] + } + } + }, + { + "pour_water_stop_pose": { + "name": "generate_affordances_from_src", + "kwargs": { + "affordance_infos": [ + { + "src_key": "pour_water_start_qpos", + "dst_key": "bottle_rotation_qpos", + "valid_funcs_name_kwargs_proc": [ + { + "name": "no_validation", + "kwargs": {}, + "pass_processes": [ + { + "name": "get_offset_qpos", + "kwargs": { + "offset_value": -75, + "joint_list_offset": [5], + "degrees": true + } + } + ] + }, + { + "name": "is_qpos_exceed", + "kwargs": { + "robot": "env.robot", + "control_part": "right_arm" + } + }, + { + "name": "is_qpos_flip", + "kwargs": { + "qpos_ref": "env.affordance_datas['right_arm_init_qpos']", + "qpos_ids": [3, 4], + "threshold": 3.455751918948773, + "mode": "delta", + "return_inverse": true + } + } + ] + }, + { + "src_key": "bottle_rotation_qpos", + "dst_key": "pour_water_stop_pose", + "valid_funcs_name_kwargs_proc": [ + { + "name": "no_validation", + "kwargs": {}, + "pass_processes": [ + { + "name": "get_fk_xpos", + "kwargs": { + "control_part": "right_arm", + "fk_func": "env.robot.compute_fk" + } + } + ] + } + ] + } + ] + } + } + }, + { + "bottle_place_pose": { + "name": "generate_affordances_from_src", + "kwargs": { + "affordance_infos": [ + { + "src_key": "bottle_grasp_pose", + "dst_key": "bottle_place_pose", + "valid_funcs_name_kwargs_proc": [ + { + "name":"no_validation", + "kwargs": {}, + "pass_processes": [ + { + "name": "get_replaced_pose", + "kwargs": { + "pose_replace_value": [0.7, -0.1], + "axis_str_replace": "xy" + } + } + ] + } + ] + }, + { + "src_key": "bottle_place_pose", + "dst_key": "bottle_place_qpos", + "valid_funcs_name_kwargs_proc": [ + { + "name": "get_ik_ret", + "kwargs": { + "ik_func": "env.robot.compute_ik", + "qpos_seed": "env.affordance_datas['pour_water_start_qpos']", + "control_part": "right_arm" + }, + "pass_processes": [ + { + "name": "get_ik_qpos", + "kwargs": { + "ik_func": "env.robot.compute_ik", + "qpos_seed": "env.affordance_datas['pour_water_start_qpos']", + "control_part": "right_arm" + } + } + ] + }, + { + "name": "is_qpos_flip", + "kwargs": { + "qpos_ref": "env.affordance_datas['right_arm_init_qpos']", + "qpos_ids": [3, 4], + "threshold": 3.455751918948773, + "mode": "delta", + "return_inverse": true + } + } + ] + } + ] + } + } + }, + { + "bottle_pre_place_qpos": { + "name": "generate_affordances_from_src", + "kwargs": { + "affordance_infos": [ + { + "src_key": "bottle_place_qpos", + "dst_key": "bottle_pre_place_qpos", + "valid_funcs_name_kwargs_proc": [ + { + "name": "no_validation", + "kwargs": {}, + "pass_processes": [ + { + "name": "get_offset_qpos", + "kwargs":{ + "offset_value": -0.05, + "joint_list_offset": [1] + } + } + ] + }, + { + "name": "is_qpos_exceed", + "kwargs": { + "robot": "env.robot", + "control_part": "right_arm" + } + }, + { + "name": "is_qpos_flip", + "kwargs": { + "qpos_ref": "env.affordance_datas['right_arm_init_qpos']", + "qpos_ids": [3, 4], + "threshold": 3.455751918948773, + "mode": "delta", + "return_inverse": true + } + } + ] + } + ] + } + } + }, + { + "compute_unoffset_for_exp":{ + "name": "compute_unoffset_for_exp", + "kwargs": { + "pose_input_output_names_changes": { + "bottle_grasp_pose": { + "output_pose_name":"bottle_grasp_pose_object_unoffset", + "pose_changes": [ + ["framechange_extrinsic_inverse", "env.affordance_datas['bottle_pose']"], + ["rotation_z_intrinsic_degrees", 90], + ["offset_-z_intrinsic", 0.025], + ["rotation_z_intrinsic_degrees", -90] + ] + }, + "bottle_place_pose": { + "output_pose_name":"bottle_place_pose_unoffset", + "pose_changes": [ + ["rotation_z_intrinsic_degrees", 90], + ["offset_-z_intrinsic", 0.025], + ["rotation_z_intrinsic_degrees", -90] + ] + }, + "pour_water_start_pose": { + "output_pose_name":"pour_water_start_pose_unoffset", + "pose_changes": [ + ["rotation_z_intrinsic_degrees", 90], + ["offset_-z_intrinsic", 0.025], + ["rotation_z_intrinsic_degrees", -90] + ] + }, + "pour_water_stop_pose": { + "output_pose_name":"pour_water_stop_pose_unoffset", + "pose_changes": [ + ["rotation_z_intrinsic_degrees", 90], + ["offset_-z_intrinsic", 0.025], + ["rotation_z_intrinsic_degrees", -90] + ] + } + } + } + } + } + ], + "left_arm": [ + { + "left_arm_init_qpos": { + "name": "generate_affordances_from_src", + "kwargs": { + "affordance_infos": [ + { + "src_key": "left_arm_init_qpos", + "dst_key": "left_arm_init_pose", + "valid_funcs_name_kwargs_proc": [ + { + "name": "no_validation", + "kwargs": {}, + "pass_processes": [ + { + "name": "get_fk_xpos", + "kwargs": { + "control_part": "left_arm", + "fk_func": "env.robot.compute_fk" + } + } + ] + } + ] + } + ] + + } + } + }, + { + "left_monitor_qpos": { + "name": "generate_affordances_from_src", + "kwargs": { + "affordance_infos": [ + { + "src_key": "left_arm_init_qpos", + "dst_key": "left_arm_monitor_qpos", + "valid_funcs_name_kwargs_proc": [ + { + "name": "no_validation", + "kwargs": {}, + "pass_processes": [ + { + "name": "get_replaced_qpos", + "kwargs": { + "replace_value": [-0.6, 1.0, -1.2, 0.0, 0.58, 0.0], + "joint_list_replace": [0, 1, 2, 3, 4, 5] + } + } + ] + } + ] + } + ] + } + } + } + ], + "left_eef": [ + { + "open": { + "name": "execute_open", + "kwargs": {} + } + }, + { + "close": { + "name": "execute_close", + "kwargs": {} + } + } + ], + "right_eef": [ + { + "ropen": { + "name": "execute_open", + "kwargs": {} + } + }, + { + "rclose": { + "name": "execute_close", + "kwargs": {} + } + } + ] + }, + "edge": { + "right_arm": [ + { + "init_to_aim": { + "src": "right_arm_init_qpos", + "sink": "right_arm_aim_qpos", + "duration": 10, + "name": "plan_trajectory", + "kwargs": { + "agent_uid": "right_arm", + "keypose_names": [ + "right_arm_init_qpos", + "right_arm_aim_qpos" + ] + } + } + }, + { + "aim_to_pre1": { + "src": "right_arm_aim_qpos", + "sink": "bottle_pre1_qpos", + "duration": 24, + "name": "plan_trajectory", + "kwargs": { + "agent_uid": "right_arm", + "keypose_names": [ + "right_arm_aim_qpos", + "bottle_pre1_qpos" + ] + } + } + }, + { + "pre1_to_grasp": { + "src": "bottle_pre1_qpos", + "sink": "bottle_grasp_qpos", + "duration": 24, + "name": "plan_trajectory", + "kwargs": { + "agent_uid": "right_arm", + "keypose_names": [ + "bottle_pre1_qpos", + "bottle_grasp_qpos" + ] + } + } + }, + { + "grasp_to_up": { + "src": "bottle_grasp_qpos", + "sink": "bottle_up_qpos", + "duration": 4, + "name": "plan_trajectory", + "kwargs": { + "agent_uid": "right_arm", + "keypose_names": [ + "bottle_grasp_qpos", + "bottle_up_qpos" + ] + } + } + }, + { + "up_to_move": { + "src": "bottle_up_qpos", + "sink": "pour_water_start_qpos", + "duration": 20, + "name": "plan_trajectory", + "kwargs": { + "agent_uid": "right_arm", + "keypose_names": [ + "bottle_up_qpos", + "pour_water_start_qpos" + ] + } + } + }, + { + "move_to_rotation": { + "src": "pour_water_start_qpos", + "sink": "bottle_rotation_qpos", + "duration": 24, + "name": "plan_trajectory", + "kwargs": { + "agent_uid": "right_arm", + "keypose_names": [ + "pour_water_start_qpos", + "bottle_rotation_qpos" + ] + } + } + }, + { + "rotation_back_to_move": { + "src": "bottle_rotation_qpos", + "sink": "pour_water_start_qpos", + "duration": 24, + "name": "plan_trajectory", + "kwargs": { + "agent_uid": "right_arm", + "keypose_names": [ + "bottle_rotation_qpos", + "pour_water_start_qpos" + ] + } + } + }, + { + "move_back_to_pre_place": { + "src": "pour_water_start_qpos", + "sink": "bottle_pre_place_qpos", + "duration": 20, + "name": "plan_trajectory", + "kwargs": { + "agent_uid": "right_arm", + "keypose_names": [ + "pour_water_start_qpos", + "bottle_pre_place_qpos" + ] + } + } + }, + { + "pre_place_back_to_place": { + "src": "bottle_pre_place_qpos", + "sink": "bottle_place_qpos", + "duration": 4, + "name": "plan_trajectory", + "kwargs": { + "agent_uid": "right_arm", + "keypose_names": [ + "bottle_pre_place_qpos", + "bottle_place_qpos" + ] + } + } + }, + { + "place_back_to_init": { + "src": "bottle_place_qpos", + "sink": "right_arm_init_qpos", + "duration": 24, + "name": "plan_trajectory", + "kwargs": { + "agent_uid": "right_arm", + "keypose_names": [ + "bottle_place_qpos", + "right_arm_init_qpos" + ] + } + } + } + ], + "left_arm": [ + { + "left_init_to_monitor": { + "src": "left_arm_init_qpos", + "sink": "left_monitor_qpos", + "duration": 15, + "name": "plan_trajectory", + "kwargs": { + "agent_uid": "left_arm", + "keypose_names": [ + "left_arm_init_qpos", + "left_arm_monitor_qpos" + ] + } + } + }, + { + "left_arm_go_back": { + "src": "left_monitor_qpos", + "sink": "left_arm_init_qpos", + "duration": 15, + "kwargs": {} + } + } + ], + "left_eef": [], + "right_eef": [ + { + "rclose0": { + "src": "ropen", + "sink": "rclose", + "duration": 11, + "name": "execute_close", + "kwargs": { + "return_action": true, + "expand": true + } + } + }, + { + "ropen0": { + "src": "rclose", + "sink": "ropen", + "duration": 11, + "name": "execute_open", + "kwargs": { + "return_action": true, + "expand": true + } + } + } + ] + }, + "sync": { + "rclose0": { + "depend_tasks": [ + "pre1_to_grasp" + ] + }, + "grasp_to_up": { + "depend_tasks": [ + "rclose0" + ] + }, + "ropen0": { + "depend_tasks": [ + "pre_place_back_to_place" + ] + }, + "place_back_to_init": { + "depend_tasks": [ + "ropen0" + ] + }, + "left_arm_go_back": { + "depend_tasks": [ + "ropen0" + ] + } + }, + "misc": { + "vis_graph": false, + "vis_gantt": false, + "warpping": true + } +} diff --git a/embodichain_tasks/configs/gym/pour_water/gym_config.json b/embodichain_tasks/configs/gym/pour_water/gym_config.json new file mode 100644 index 00000000..6668f4b0 --- /dev/null +++ b/embodichain_tasks/configs/gym/pour_water/gym_config.json @@ -0,0 +1,351 @@ +{ + "id": "PourWater-v3", + "max_episodes": 5, + "max_episode_steps": 300, + "env": { + "events": { + "record_camera": { + "func": "record_camera_data", + "mode": "interval", + "interval_step": 1, + "params": { + "name": "cam1", + "resolution": [320, 240], + "eye": [2, 0, 2], + "target": [0.5, 0, 1] + } + }, + "random_light": { + "func": "randomize_light", + "mode": "interval", + "interval_step": 10, + "params": { + "entity_cfg": {"uid": "light_1"}, + "position_range": [[-0.5, -0.5, 2], [0.5, 0.5, 2]], + "color_range": [[0.6, 0.6, 0.6], [1, 1, 1]], + "intensity_range": [50.0, 100.0] + } + }, + "init_bottle_pose": { + "func": "randomize_rigid_object_pose", + "mode": "reset", + "params": { + "entity_cfg": {"uid": "bottle"}, + "position_range": [[-0.08, -0.12, 0.0], [0.08, 0.04, 0.0]], + "relative_position": true + } + }, + "init_cup_pose": { + "func": "randomize_rigid_object_pose", + "mode": "reset", + "params": { + "entity_cfg": {"uid": "cup"}, + "position_range": [[-0.08, -0.04, 0.0], [0.08, 0.12, 0.0]], + "relative_position": true + } + }, + "prepare_extra_attr": { + "func": "prepare_extra_attr", + "mode": "reset", + "params": { + "attrs": [ + { + "name": "object_lengths", + "mode": "callable", + "entity_uids": "all_objects", + "func_name": "compute_object_length", + "func_kwargs": { + "is_svd_frame": true, + "sample_points": 5000 + } + }, + { + "name": "grasp_pose_object", + "mode": "static", + "entity_cfg": { + "uid": "bottle" + }, + "value": [[ + [0.32243, 0.03245, 0.94604, 0.025], + [0.00706, -0.99947, 0.03188, -0.0 ], + [0.94657, -0.0036 , -0.32249, 0.0 ], + [0.0 , 0.0 , 0.0 , 1.0 ] + ]] + }, + { + "name": "left_arm_base_pose", + "mode": "callable", + "entity_cfg": { + "uid": "CobotMagic" + }, + "func_name": "get_link_pose", + "func_kwargs": { + "link_name": "left_arm_base", + "to_matrix": true + } + }, + { + "name": "right_arm_base_pose", + "mode": "callable", + "entity_cfg": { + "uid": "CobotMagic" + }, + "func_name": "get_link_pose", + "func_kwargs": { + "link_name": "right_arm_base", + "to_matrix": true + } + } + ] + } + }, + "register_info_to_env": { + "func": "register_info_to_env", + "mode": "reset", + "params": { + "registry": [ + { + "entity_cfg": { + "uid": "bottle" + }, + "pose_register_params": { + "compute_relative": false, + "compute_pose_object_to_arena": true, + "to_matrix": true + } + }, + { + "entity_cfg": { + "uid": "cup" + }, + "pose_register_params": { + "compute_relative": false, + "compute_pose_object_to_arena": true, + "to_matrix": true + } + }, + { + "entity_cfg": { + "uid": "CobotMagic", + "control_parts": ["left_arm"] + }, + "attrs": ["left_arm_base_pose"], + "pose_register_params": { + "compute_relative": "cup", + "compute_pose_object_to_arena": false, + "to_matrix": true + }, + "prefix": false + }, + { + "entity_cfg": { + "uid": "CobotMagic", + "control_parts": ["right_arm"] + }, + "attrs": ["right_arm_base_pose"], + "pose_register_params": { + "compute_relative": "bottle", + "compute_pose_object_to_arena": false, + "to_matrix": true + }, + "prefix": false + } + ], + "registration": "affordance_datas", + "sim_update": true + } + }, + "random_material": { + "func": "randomize_visual_material", + "mode": "interval", + "interval_step": 10, + "params": { + "entity_cfg": {"uid": "table"}, + "random_texture_prob": 0.5, + "texture_path": "CocoBackground/coco", + "base_color_range": [[0.2, 0.2, 0.2], [1.0, 1.0, 1.0]] + } + }, + "random_cup_material": { + "func": "randomize_visual_material", + "mode": "interval", + "interval_step": 10, + "params": { + "entity_cfg": {"uid": "cup"}, + "random_texture_prob": 0.5, + "texture_path": "CocoBackground/coco", + "base_color_range": [[0.2, 0.2, 0.2], [1.0, 1.0, 1.0]] + } + }, + "random_bottle_material": { + "func": "randomize_visual_material", + "mode": "interval", + "interval_step": 10, + "params": { + "entity_cfg": {"uid": "bottle"}, + "random_texture_prob": 0.5, + "texture_path": "CocoBackground/coco", + "base_color_range": [[0.2, 0.2, 0.2], [1.0, 1.0, 1.0]] + } + }, + "random_robot_init_eef_pose": { + "func": "randomize_robot_eef_pose", + "mode": "reset", + "params": { + "entity_cfg": {"uid": "CobotMagic", "control_parts": ["left_arm", "right_arm"]}, + "position_range": [[-0.01, -0.01, -0.01], [0.01, 0.01, 0]] + } + } + }, + "observations": { + "norm_robot_eef_joint": { + "func": "normalize_robot_joint_data", + "mode": "modify", + "name": "robot/qpos", + "params": { + "joint_ids": [6, 13] + } + } + }, + "dataset": { + "lerobot": { + "func": "LeRobotRecorder", + "mode": "save", + "params": { + "robot_meta": { + "robot_type": "CobotMagic", + "control_freq": 25 + }, + "instruction": { + "lang": "Pour water from bottle to cup" + }, + "extra": { + "scene_type": "Commercial", + "task_description": "Pour water", + "data_type": "sim" + }, + "use_videos": true + } + } + }, + "control_parts": ["left_arm", "left_eef", "right_arm", "right_eef"] + }, + "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": 960, + "height": 540, + "intrinsics": [488.1665344238281, 488.1665344238281, 480, 270], + "extrinsics": { + "eye": [0.35368482807598, 0.014695524383058989, 1.4517046071614774], + "target": [0.8586357573287919, 0, 0.5232553674540066], + "up": [0.9306678549330372, -0.0005600064212467153, 0.3658647703553347] + } + }, + { + "sensor_type": "Camera", + "uid": "cam_right_wrist", + "width": 640, + "height": 480, + "intrinsics": [488.1665344238281, 488.1665344238281, 322.7323303222656, 213.17434692382812], + "extrinsics": { + "parent": "right_link6", + "pos": [-0.08, 0.0, 0.04], + "quat": [0.15304635, 0.69034543, -0.69034543, -0.15304635] + } + }, + { + "sensor_type": "Camera", + "uid": "cam_left_wrist", + "width": 640, + "height": 480, + "intrinsics": [488.1665344238281, 488.1665344238281, 322.7323303222656, 213.17434692382812], + "extrinsics": { + "parent": "left_link6", + "pos": [-0.08, 0.0, 0.04], + "quat": [0.15304635, 0.69034543, -0.69034543, -0.15304635] + } + } + ], + "light": { + "direct": [ + { + "uid": "light_1", + "light_type": "point", + "color": [1.0, 1.0, 1.0], + "intensity": 50.0, + "init_pos": [2, 0, 2], + "radius": 10.0 + } + ] + }, + "background": [ + { + "uid": "table", + "shape": { + "shape_type": "Mesh", + "fpath": "CircleTableSimple/circle_table_simple.ply", + "compute_uv": true + }, + "attrs" : { + "mass": 10.0, + "static_friction": 0.95, + "dynamic_friction": 0.9, + "restitution": 0.01 + }, + "body_scale": [1, 1, 1], + "body_type": "kinematic", + "init_pos": [0.725, 0.0, 0.825], + "init_rot": [0, 90, 0] + } + ], + "rigid_object": [ + { + "uid":"cup", + "shape": { + "shape_type": "Mesh", + "fpath": "PaperCup/paper_cup.ply", + "compute_uv": true + }, + "attrs" : { + "mass": 0.01, + "contact_offset": 0.003, + "rest_offset": 0.001, + "restitution": 0.01, + "max_depenetration_velocity": 1e1, + "min_position_iters": 32, + "min_velocity_iters":8 + }, + "init_pos": [0.75, 0.1, 0.9], + "body_scale":[0.75, 0.75, 1.0], + "max_convex_hull_num": 8 + }, + { + "uid":"bottle", + "shape": { + "shape_type": "Mesh", + "fpath": "ScannedBottle/kashijia_processed.ply", + "compute_uv": true + }, + "attrs" : { + "mass": 0.01, + "contact_offset": 0.003, + "rest_offset": 0.001, + "restitution": 0.01, + "max_depenetration_velocity": 1e1, + "min_position_iters": 32, + "min_velocity_iters":8 + }, + "init_pos": [0.75, -0.1, 0.932], + "body_scale":[1, 1, 1], + "max_convex_hull_num": 8 + } + ] +} \ No newline at end of file diff --git a/embodichain_tasks/configs/gym/scoop_ice/gym_config.json b/embodichain_tasks/configs/gym/scoop_ice/gym_config.json new file mode 100644 index 00000000..331f8db6 --- /dev/null +++ b/embodichain_tasks/configs/gym/scoop_ice/gym_config.json @@ -0,0 +1,235 @@ +{ + "id": "ScoopIce-v1", + "max_episodes": 5, + "env": { + "events": { + "drop_ice":{ + "func": "drop_rigid_object_group_sequentially", + "mode": "reset", + "params": { + "entity_cfg": {"uid": "ice_cubes"}, + "drop_position": [0.5, -0.05, 1.0], + "position_range": [[-0.12, -0.12, 0], [0.12, 0.12, 0]], + "physics_step": 10 + } + }, + "init_scoop_pose": { + "func": "randomize_rigid_object_pose", + "mode": "reset", + "params": { + "entity_cfg": {"uid": "scoop"}, + "position_range": [[0.45, -0.27, 1.05], [0.45, -0.27, 1.05]], + "rotation_range": [[30, -25, 180], [30, -25, 180]], + "relative_position": false, + "relative_rotation": false + } + }, + "init_cup_pose": { + "func": "randomize_rigid_object_pose", + "mode": "reset", + "params": { + "entity_cfg": {"uid": "paper_cup"}, + "position_range": [[0.455, 0.112, 1.05], [0.455, 0.112, 1.05]], + "rotation_range": [[0, 0, 0], [0, 0, 0]], + "relative_position": false, + "relative_rotation": false + } + }, + "record_camera": { + "func": "record_camera_data", + "mode": "interval", + "interval_step": 1, + "params": { + "name": "cam1", + "resolution": [320, 240], + "eye": [0, -1, 2], + "target": [0.5, 0, 1] + } + } + } + }, + "robot": { + "robot_type": "DexforceW1", + "init_pos": [0.0, 0.0, 0], + "init_qpos":[ + 0.42241, -1.11061, 0.55116, 0.01815, 0.00002, -0.43273, + -0.30339, 0.2412 , -1.20074, 0.72621, 0.40264, -0.41044, + -1.34341, 1.27664, -0.42869, -0.30873, -0.23608, 0.54272, + -0.12095, -0.16959, 0.00011, 0.00002, 0.00009, 0.00003, + 1.50006, 0.30045, 0.30028, 0.30051, 0.30037, 0.90003, + 0.23262, 0.24298, 0.22003, 0.19901, -0 , 0.59305, + 0.6033 , 0.58056, 0.55942, 0.29999 + ], + "solver_cfg": { + "left_arm": { + "class_type": "PytorchSolver", + "end_link_name": "left_ee", + "root_link_name": "left_arm_base", + "tcp": + [ + [-1.0, 0.0, 0.0, 0.012], + [0.0, 0.0, 1.0, 0.0675], + [0.0, 1.0, 0.0, 0.127], + [0.0, 0.0, 0.0, 1.0] + ] + }, + "right_arm": { + "class_type": "PytorchSolver", + "end_link_name": "right_ee", + "root_link_name": "right_arm_base", + "tcp": + [ + [1.0, 0.0, 0.0, 0.012], + [0.0, 0.0, -1.0, -0.0675], + [0.0, 1.0, 0.0, 0.127], + [0.0, 0.0, 0.0, 1.0] + ] + } + } + }, + "sensor": [ + { + "sensor_type": "StereoCamera", + "uid": "cam_high", + "width": 960, + "height": 540, + "enable_mask": true, + "enable_depth": true, + "left_to_right_pos": [0.059684025824163614, 0, 0], + "intrinsics": [453.851402686215, 453.8347628855552, 469.827725021235, 258.6656181845155], + "intrinsics_right": [453.4536601653505, 453.3306024582175, 499.13697412367776, 297.7176248477935], + "extrinsics": { + "parent": "eyes" + } + }, + { + "sensor_type": "Camera", + "uid": "cam_right_wrist", + "width": 640, + "height": 360, + "enable_mask": true, + "intrinsics": [337.0, 325.0, 320.0, 180.0], + "extrinsics": { + "parent": "right_ee", + "pos": [0.09, 0.05, 0.04], + "quat": [0.36497168, -0.11507513, 0.88111957, 0.27781593] + } + }, + { + "sensor_type": "Camera", + "uid": "cam_left_wrist", + "width": 640, + "height": 360, + "enable_mask": true, + "intrinsics": [337.0, 325.0, 320.0, 180.0], + "extrinsics": { + "parent": "left_ee", + "pos": [0.09, -0.05, 0.04], + "quat": [0.27781593, 0.88111957, -0.11507513, 0.36497168] + } + } + ], + "light": { + }, + "background": [ + { + "uid": "table", + "shape": { + "shape_type": "Mesh", + "fpath": "CircleTableSimple/circle_table_simple.ply" + }, + "attrs" : { + "mass": 1.0, + "static_friction": 0.95, + "dynamic_friction": 0.9, + "restitution": 0.05 + }, + "body_type": "kinematic", + "init_pos": [0.80, 0, 0.54], + "init_rot": [0, 90, 0] + } + ], + "rigid_object": [ + { + "uid": "scoop", + "shape": { + "shape_type": "Mesh", + "fpath": "ScoopIceNewEnv/scoop.ply" + }, + "attrs" : { + "mass": 0.5, + "static_friction": 0.95, + "dynamic_friction": 0.9, + "restitution": 0.0, + "min_position_iters": 32, + "min_velocity_iters": 8 + }, + "max_convex_hull_num": 8, + "init_pos": [0, 10, 10] + }, + { + "uid": "paper_cup", + "shape": { + "shape_type": "Mesh", + "fpath": "PaperCup/paper_cup.ply" + }, + "attrs" : { + "mass": 0.5, + "static_friction": 0.95, + "dynamic_friction": 0.9, + "restitution": 0.0, + "min_position_iters": 32, + "min_velocity_iters": 8 + }, + "max_convex_hull_num": 16, + "init_pos": [0, 10, 10] + } + ], + "rigid_object_group": [ + { + "uid": "ice_cubes", + "max_num": 300, + "folder_path": "ScoopIceNewEnv/ice_mesh_small", + "ext": ".obj", + "rigid_objects": { + "obj": { + "attrs" : { + "mass": 0.004, + "contact_offset": 0.001, + "rest_offset": 0, + "dynamic_friction": 0.05, + "static_friction": 0.1, + "restitution": 0.00, + "min_position_iters": 32, + "min_velocity_iters": 8, + "max_depenetration_velocity": 1.0 + }, + "shape": { + "shape_type": "Mesh" + }, + "init_pos": [0, 0, 2], + "body_scale": [1.0, 1.0, 1.0] + } + } + } + ], + "articulation": [ + { + "uid": "container", + "fpath": "ScoopIceNewEnv/IceContainer/ice_container.urdf", + "init_pos": [0.635, -0.04, 0.94], + "init_rot": [0, 0, -80], + "attrs": { + "mass": 1.0, + "dynamic_friction": 0.05, + "static_friction": 0.1, + "max_depenetration_velocity": 1.0 + }, + "drive_pros": { + "stiffness": 1.0, + "damping": 0.1, + "max_effort": 100.0 + } + } + ] +} \ No newline at end of file diff --git a/embodichain_tasks/configs/gym/special/simple_task_ur10.json b/embodichain_tasks/configs/gym/special/simple_task_ur10.json new file mode 100644 index 00000000..1c722b1a --- /dev/null +++ b/embodichain_tasks/configs/gym/special/simple_task_ur10.json @@ -0,0 +1,104 @@ +{ + "id": "SimpleTask-v1", + "max_episodes": 24, + "env": { + "events": { + "random_light": { + "func": "randomize_light", + "mode": "interval", + "interval_step": 20, + "params": { + "entity_cfg": {"uid": "light_1"}, + "position_range": [[-0.5, -0.5, 2], [0.5, 0.5, 2]], + "color_range": [[0.6, 0.6, 0.6], [1, 1, 1]], + "intensity_range": [50.0, 100.0] + } + }, + "random_material": { + "func": "randomize_visual_material", + "mode": "interval", + "interval_step": 50, + "params": { + "entity_cfg": {"uid": "table"}, + "random_texture_prob": 0.0, + "texture_path": "CocoBackground/coco", + "base_color_range": [[0.2, 0.2, 0.2], [1.0, 1.0, 1.0]] + } + } + }, + "dataset": { + "lerobot": { + "func": "LeRobotRecorder", + "mode": "save", + "params": { + "robot_meta": { + "robot_type": "UR10", + "control_freq": 3 + }, + "instruction": { + "lang": "Acting with Oscillatory motion" + }, + "extra": { + "scene_type": "commercial", + "task_description": "Oscillatory motion", + "data_type": "sim" + }, + "use_videos": false + } + } + } + }, + "robot": { + "uid": "UR10", + "fpath": "UniversalRobots/UR10/UR10.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, 3], + "target": [0, 0, 1] + } + } + ], + "light": { + "direct": [ + { + "uid": "light_1", + "light_type": "point", + "color": [1.0, 1.0, 1.0], + "intensity": 50.0, + "init_pos": [2, 0, 2], + "radius": 10.0 + } + ] + }, + "background": [ + { + "uid": "table", + "shape": { + "shape_type": "Mesh", + "fpath": "CircleTableSimple/circle_table_simple.ply", + "compute_uv": true + }, + "attrs" : { + "mass": 10.0, + "static_friction": 0.95, + "dynamic_friction": 0.9, + "restitution": 0.01 + }, + "body_scale": [1, 1, 1], + "body_type": "kinematic", + "init_pos": [0.8, 0.0, 0.825], + "init_rot": [0, 90, 0] + } + ], + "rigid_object": [ + ] +} \ No newline at end of file diff --git a/embodichain_tasks/configs/gym/stack_blocks_two/cobot_magic_3cam.json b/embodichain_tasks/configs/gym/stack_blocks_two/cobot_magic_3cam.json new file mode 100644 index 00000000..460a53c2 --- /dev/null +++ b/embodichain_tasks/configs/gym/stack_blocks_two/cobot_magic_3cam.json @@ -0,0 +1,191 @@ +{ + "id": "StackBlocksTwo-v1", + "max_episodes": 5, + "env": { + "events": { + "random_light": { + "func": "randomize_light", + "mode": "interval", + "interval_step": 10, + "params": { + "entity_cfg": {"uid": "light_1"}, + "position_range": [[-0.5, -0.5, 2], [0.5, 0.5, 2]], + "color_range": [[0.6, 0.6, 0.6], [1, 1, 1]], + "intensity_range": [50.0, 100.0] + } + }, + "init_block_1_pose": { + "func": "randomize_rigid_object_pose", + "mode": "reset", + "params": { + "entity_cfg": {"uid": "block_1"}, + "position_range": [[-0.08, -0.08, 0.0], [0.08, 0.08, 0.0]], + "rotation_range": [[0, 0, -0.75], [0, 0, 0.75]], + "relative_position": true + } + }, + "init_block_2_pose": { + "func": "randomize_rigid_object_pose", + "mode": "reset", + "params": { + "entity_cfg": {"uid": "block_2"}, + "position_range": [[-0.08, -0.08, 0.0], [0.08, 0.08, 0.0]], + "rotation_range": [[0, 0, -0.75], [0, 0, 0.75]], + "relative_position": true + } + } + }, + "observations": { + "norm_robot_eef_joint": { + "func": "normalize_robot_joint_data", + "mode": "modify", + "name": "robot/qpos", + "params": { + "joint_ids": [6, 13] + } + }, + "block_1_pose": { + "func": "get_rigid_object_pose", + "mode": "add", + "name": "block_1_pose", + "params": { + "entity_cfg": {"uid": "block_1"} + } + }, + "block_2_pose": { + "func": "get_rigid_object_pose", + "mode": "add", + "name": "block_2_pose", + "params": { + "entity_cfg": {"uid": "block_2"} + } + } + } + }, + "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": "StereoCamera", + "uid": "cam_high", + "width": 960, + "height": 540, + "enable_mask": true, + "enable_depth": true, + "left_to_right_pos": [0.059684025824163614, 0, 0], + "intrinsics": [453.851402686215, 453.8347628855552, 469.827725021235, 258.6656181845155], + "intrinsics_right": [453.4536601653505, 453.3306024582175, 499.13697412367776, 297.7176248477935], + "extrinsics": { + "eye": [0.35368482807598, 0.014695524383058989, 1.4517046071614774], + "target": [0.7186357573287919, -0.054534732904795505, 0.5232553674540066], + "up": [0.9306678549330372, -0.0005600064212467153, 0.3658647703553347] + } + }, + { + "sensor_type": "Camera", + "uid": "cam_right_wrist", + "width": 640, + "height": 480, + "enable_mask": true, + "intrinsics": [488.1665344238281, 488.1665344238281, 322.7323303222656, 213.17434692382812], + "extrinsics": { + "parent": "right_link6", + "pos": [-0.08, 0.0, 0.04], + "quat": [0.15304635, 0.69034543, -0.69034543, -0.15304635] + } + }, + { + "sensor_type": "Camera", + "uid": "cam_left_wrist", + "width": 640, + "height": 480, + "enable_mask": true, + "intrinsics": [488.1665344238281, 488.1665344238281, 322.7323303222656, 213.17434692382812], + "extrinsics": { + "parent": "left_link6", + "pos": [-0.08, 0.0, 0.04], + "quat": [0.15304635, 0.69034543, -0.69034543, -0.15304635] + } + } + ], + "light": { + "direct": [ + { + "uid": "light_1", + "light_type": "point", + "color": [1.0, 1.0, 1.0], + "intensity": 50.0, + "init_pos": [2, 0, 2], + "radius": 10.0 + } + ] + }, + "background": [ + { + "uid": "table", + "shape": { + "shape_type": "Mesh", + "fpath": "CircleTableSimple/circle_table_simple.ply" + }, + "attrs" : { + "mass": 10.0, + "static_friction": 0.95, + "dynamic_friction": 0.9, + "restitution": 0.01 + }, + "body_scale": [1, 1, 1], + "body_type": "kinematic", + "init_pos": [0.725, 0.0, 0.825], + "init_rot": [0, 90, 0] + } + ], + "rigid_object": [ + { + "uid":"block_1", + "shape": { + "shape_type": "Cube", + "size": [0.05, 0.05, 0.05] + }, + "attrs" : { + "mass": 0.05, + "static_friction": 0.5, + "dynamic_friction": 0.5, + "restitution": 0.0, + "contact_offset": 0.003, + "rest_offset": 0.001, + "max_depenetration_velocity": 1e1, + "min_position_iters": 32, + "min_velocity_iters": 8 + }, + "init_pos": [0.75, -0.1, 0.9], + "body_scale":[1, 1, 1], + "max_convex_hull_num": 1 + }, + { + "uid":"block_2", + "shape": { + "shape_type": "Cube", + "size": [0.05, 0.05, 0.05] + }, + "attrs" : { + "mass": 0.05, + "static_friction": 0.5, + "dynamic_friction": 0.5, + "restitution": 0.0, + "contact_offset": 0.003, + "rest_offset": 0.001, + "max_depenetration_velocity": 1e1, + "min_position_iters": 32, + "min_velocity_iters": 8 + }, + "init_pos": [0.75, 0.1, 0.9], + "body_scale":[1, 1, 1], + "max_convex_hull_num": 1 + } + ] +} + diff --git a/embodichain_tasks/configs/gym/stack_cups/cobot_magic_3cam.json b/embodichain_tasks/configs/gym/stack_cups/cobot_magic_3cam.json new file mode 100644 index 00000000..09daa149 --- /dev/null +++ b/embodichain_tasks/configs/gym/stack_cups/cobot_magic_3cam.json @@ -0,0 +1,202 @@ +{ + "id": "StackCups-v1", + "max_episodes": 5, + "env": { + "events": { + "random_light": { + "func": "randomize_light", + "mode": "interval", + "interval_step": 10, + "params": { + "entity_cfg": {"uid": "light_1"}, + "position_range": [[-0.5, -0.5, 2], [0.5, 0.5, 2]], + "color_range": [[0.6, 0.6, 0.6], [1, 1, 1]], + "intensity_range": [50.0, 100.0] + } + }, + "init_cup_1_pose": { + "func": "randomize_rigid_object_pose", + "mode": "reset", + "params": { + "entity_cfg": {"uid": "cup_1"}, + "position_range": [[0.66, -0.15, 0.86], [0.74, -0.05, 0.86]], + "rotation_range": [[0, 0, -0.52], [0, 0, 0.52]], + "relative_position": false + } + }, + "init_cup_2_pose": { + "func": "randomize_rigid_object_pose", + "mode": "reset", + "params": { + "entity_cfg": {"uid": "cup_2"}, + "position_range": [[0.76, -0.15, 0.86], [0.84, -0.05, 0.86]], + "rotation_range": [[0, 0, -0.52], [0, 0, 0.52]], + "relative_position": false + } + } + }, + "observations": { + "norm_robot_eef_joint": { + "func": "normalize_robot_joint_data", + "mode": "modify", + "name": "robot/qpos", + "params": { + "joint_ids": [6, 13] + } + }, + "cup_1_pose": { + "func": "get_rigid_object_pose", + "mode": "add", + "name": "cup_1_pose", + "params": { + "entity_cfg": {"uid": "cup_1"} + } + }, + "cup_2_pose": { + "func": "get_rigid_object_pose", + "mode": "add", + "name": "cup_2_pose", + "params": { + "entity_cfg": {"uid": "cup_2"} + } + } + } + }, + "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": "StereoCamera", + "uid": "cam_high", + "width": 960, + "height": 540, + "enable_mask": true, + "enable_depth": true, + "left_to_right_pos": [0.059684025824163614, 0, 0], + "intrinsics": [453.851402686215, 453.8347628855552, 469.827725021235, 258.6656181845155], + "intrinsics_right": [453.4536601653505, 453.3306024582175, 499.13697412367776, 297.7176248477935], + "extrinsics": { + "eye": [0.35368482807598, 0.014695524383058989, 1.4517046071614774], + "target": [0.7186357573287919, -0.054534732904795505, 0.5232553674540066], + "up": [0.9306678549330372, -0.0005600064212467153, 0.3658647703553347] + } + }, + { + "sensor_type": "Camera", + "uid": "cam_right_wrist", + "width": 640, + "height": 480, + "enable_mask": true, + "intrinsics": [488.1665344238281, 488.1665344238281, 322.7323303222656, 213.17434692382812], + "extrinsics": { + "parent": "right_link6", + "pos": [-0.08, 0.0, 0.04], + "quat": [0.15304635, 0.69034543, -0.69034543, -0.15304635] + } + }, + { + "sensor_type": "Camera", + "uid": "cam_left_wrist", + "width": 640, + "height": 480, + "enable_mask": true, + "intrinsics": [488.1665344238281, 488.1665344238281, 322.7323303222656, 213.17434692382812], + "extrinsics": { + "parent": "left_link6", + "pos": [-0.08, 0.0, 0.04], + "quat": [0.15304635, 0.69034543, -0.69034543, -0.15304635] + } + } + ], + "light": { + "direct": [ + { + "uid": "light_1", + "light_type": "point", + "color": [1.0, 1.0, 1.0], + "intensity": 50.0, + "init_pos": [2, 0, 2], + "radius": 10.0 + } + ] + }, + "background": [ + { + "uid": "table", + "shape": { + "shape_type": "Mesh", + "fpath": "CircleTableSimple/circle_table_simple.ply" + }, + "attrs" : { + "mass": 10.0, + "static_friction": 0.95, + "dynamic_friction": 0.9, + "restitution": 0.01 + }, + "body_scale": [1, 1, 1], + "body_type": "kinematic", + "init_pos": [0.725, 0.0, 0.825], + "init_rot": [0, 90, 0] + } + ], + "rigid_object": [ + { + "uid":"cup_1", + "shape": { + "shape_type": "Mesh", + "fpath": "PaperCup/paper_cup.ply" + }, + "attrs" : { + "mass": 0.01, + "static_friction": 1.0, + "dynamic_friction": 1.0, + "restitution": 0.0, + "contact_offset": 0.003, + "rest_offset": 0.001, + "max_depenetration_velocity": 2.0, + "linear_damping": 2.0, + "angular_damping": 2.0, + "max_linear_velocity": 5.0, + "max_angular_velocity": 10.0, + "min_position_iters": 32, + "min_velocity_iters": 8 + }, + "init_pos": [0.70, -0.1, 0.86], + "init_rot": [0, 0, 0], + "body_scale":[0.75, 0.75, 1.0], + "max_convex_hull_num": 8 + }, + { + "uid":"cup_2", + "shape": { + "shape_type": "Mesh", + "fpath": "PaperCup/paper_cup.ply" + }, + "attrs" : { + "mass": 0.01, + "static_friction": 1.0, + "dynamic_friction": 1.0, + "restitution": 0.0, + "contact_offset": 0.003, + "rest_offset": 0.001, + "max_depenetration_velocity": 2.0, + "linear_damping": 2.0, + "angular_damping": 2.0, + "max_linear_velocity": 5.0, + "max_angular_velocity": 10.0, + "min_position_iters": 32, + "min_velocity_iters": 8 + }, + "init_pos": [0.80, -0.1, 0.86], + "init_rot": [0, 0, 0], + "body_scale":[0.75, 0.75, 1.0], + "max_convex_hull_num": 8 + } + ] +} + + diff --git a/embodichain_tasks/embodichain_tasks/__init__.py b/embodichain_tasks/embodichain_tasks/__init__.py new file mode 100644 index 00000000..a01e66a2 --- /dev/null +++ b/embodichain_tasks/embodichain_tasks/__init__.py @@ -0,0 +1,31 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Official task environments for EmbodiChain. + +Importing this package triggers auto-registration of all task environments +via recursive sub-package import. Each task sub-package's ``__init__.py`` +calls ``@register_env`` which registers the environment in gymnasium's +global registry. +""" + +from __future__ import annotations + +from .utils.importer import import_packages + +_BLACKLIST = ["utils"] + +import_packages(__name__, _BLACKLIST) diff --git a/embodichain_tasks/embodichain_tasks/rl/__init__.py b/embodichain_tasks/embodichain_tasks/rl/__init__.py new file mode 100644 index 00000000..e1c548e0 --- /dev/null +++ b/embodichain_tasks/embodichain_tasks/rl/__init__.py @@ -0,0 +1,32 @@ +# ---------------------------------------------------------------------------- +# 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 copy import deepcopy +from embodichain.lab.gym.utils import registration as env_registry +from embodichain.lab.gym.envs.embodied_env import EmbodiedEnvCfg + + +def build_env(env_id: str, base_env_cfg: EmbodiedEnvCfg): + """Create env from registry id, auto-inferring cfg class (EnvName -> EnvNameCfg).""" + env = env_registry.make(env_id, cfg=deepcopy(base_env_cfg)) + return env + + +__all__ = [ + "build_env", +] diff --git a/embodichain_tasks/embodichain_tasks/rl/basic/__init__.py b/embodichain_tasks/embodichain_tasks/rl/basic/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/embodichain_tasks/embodichain_tasks/rl/basic/cart_pole.py b/embodichain_tasks/embodichain_tasks/rl/basic/cart_pole.py new file mode 100644 index 00000000..bac1e7f2 --- /dev/null +++ b/embodichain_tasks/embodichain_tasks/rl/basic/cart_pole.py @@ -0,0 +1,78 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +import torch +from typing import Dict, Any, 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("CartPoleRL", max_episode_steps=50, override=True) +class CartPoleEnv(EmbodiedEnv): + """ + CartPole balancing task for reinforcement learning. + + The agent controls a cart (robot hand joint) to keep a pole balanced near the upright + position by regulating its angle and angular velocity. Episodes are considered + successful when the pole remains close to vertical with low velocity, and they + terminate either when a maximum number of steps is reached or when the pole falls + beyond an allowed tilt threshold. + """ + + def __init__(self, cfg=None, **kwargs): + if cfg is None: + cfg = EmbodiedEnvCfg() + super().__init__(cfg, **kwargs) + + def get_reward(self, obs, action, info): + """Get the reward for the current step (pole upward reward). + + Each SimulationManager env must implement its own get_reward function to define the reward function for the task, If the + env is considered for RL/IL training. + + Args: + obs: The observation from the environment. + action: The action applied to the robot agent. + info: The info dictionary. + + Returns: + The reward for the current step. + """ + pole_qpos = self.robot.get_qpos(name="hand").reshape(-1) # [num_envs, ] + + normalized_upward = torch.abs(pole_qpos) / torch.pi + reward = 1.0 - normalized_upward + return reward + + def compute_task_state( + self, **kwargs + ) -> Tuple[torch.Tensor, torch.Tensor, Dict[str, Any]]: + qpos = self.robot.get_qpos(name="hand").reshape(-1) # [num_envs, ] + qvel = self.robot.get_qvel(name="hand").reshape(-1) # [num_envs, ] + upward_distance = torch.abs(qpos) + balance = torch.logical_and(upward_distance < 0.02, torch.abs(qvel) < 0.05) + at_final_step = self._elapsed_steps >= self.max_episode_steps - 1 + is_success = torch.logical_and(at_final_step, balance) + is_fail = torch.zeros(self.num_envs, device=self.device, dtype=torch.bool) + metrics = {"distance_to_goal": upward_distance} + return is_success, is_fail, metrics + + def check_truncated(self, obs: EnvObs, info: Dict[str, Any]) -> torch.Tensor: + pole_qpos = self.robot.get_qpos(name="hand").reshape(-1) + is_fallen = torch.abs(pole_qpos) > torch.pi * 0.5 + return is_fallen diff --git a/embodichain_tasks/embodichain_tasks/rl/push_cube.py b/embodichain_tasks/embodichain_tasks/rl/push_cube.py new file mode 100644 index 00000000..361ee3ac --- /dev/null +++ b/embodichain_tasks/embodichain_tasks/rl/push_cube.py @@ -0,0 +1,65 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +import torch +from typing import Dict, Any, 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("PushCubeRL", max_episode_steps=50, override=True) +class PushCubeEnv(EmbodiedEnv): + """Push cube task for reinforcement learning. + + The task involves pushing a cube to a target goal position using a robotic arm. + The reward consists of reaching reward, placing reward, action penalty, and success bonus. + """ + + def __init__(self, cfg=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]]: + cube = self.sim.get_rigid_object("cube") + cube_pos = cube.get_local_pose(to_matrix=True)[:, :3, 3] + + # Check if goal_pose is defined (set by randomize_target_pose event) + goal_pose = getattr(self, "goal_pose", None) + if goal_pose is not None: + goal_pos = goal_pose[:, :3, 3] + xy_distance = torch.norm(cube_pos[:, :2] - goal_pos[:, :2], dim=1) + is_success = xy_distance < self.success_threshold + else: + xy_distance = torch.zeros(self.num_envs, device=self.device) + 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 = {"distance_to_goal": xy_distance} + + return is_success, is_fail, metrics + + def check_truncated(self, obs: EnvObs, info: Dict[str, Any]) -> torch.Tensor: + cube = self.sim.get_rigid_object("cube") + cube_pos = cube.get_local_pose(to_matrix=True)[:, :3, 3] + is_fallen = cube_pos[:, 2] < -0.1 + return is_fallen diff --git a/embodichain_tasks/embodichain_tasks/special/__init__.py b/embodichain_tasks/embodichain_tasks/special/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/embodichain_tasks/embodichain_tasks/special/simple_task.py b/embodichain_tasks/embodichain_tasks/special/simple_task.py new file mode 100644 index 00000000..97c50731 --- /dev/null +++ b/embodichain_tasks/embodichain_tasks/special/simple_task.py @@ -0,0 +1,87 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +import torch + +from embodichain.lab.gym.envs import EmbodiedEnv, EmbodiedEnvCfg +from embodichain.lab.gym.utils.registration import register_env +from embodichain.utils import logger + +__all__ = ["SimpleTaskEnv"] + + +@register_env("SimpleTask-v1", max_episode_steps=600) +class SimpleTaskEnv(EmbodiedEnv): + """A demo environment with sinusoidal trajectory + + Args: + EmbodiedEnv (_type_): _description_ + """ + + def __init__(self, cfg: EmbodiedEnvCfg = None, **kwargs): + super().__init__(cfg, **kwargs) + + def create_demo_action_list(self, *args, **kwargs): + """ + Create a demonstration action list for the current task. + + This demo creates a simple sinusoidal trajectory for the robot joints. + + Returns: + list: A list of demo actions generated by the task. + """ + action_list = [] + num_steps = 100 + + # Get initial pose + init_pose = self.robot.get_qpos() # shape: (num_envs, num_joints) + + # Create a sinusoidal trajectory + for i in range(num_steps): + # Calculate phase for sinusoidal motion + t = i / num_steps # 0 to 1 + phase = torch.full( + (init_pose.shape[0],), t * 2 * torch.pi, device=self.device + ) # repeat for num_envs + + # Create sinusoidal offsets for each joint + # Joint 0: horizontal movement + # Joint 1: vertical movement + # Other joints: smaller oscillations + offset = torch.zeros_like( + init_pose, dtype=torch.float32, device=self.device + ) + offset[:, 0] = torch.sin(phase) * 0.3 # ±0.3 rad + offset[:, 1] = torch.cos(phase) * 0.2 # ±0.2 rad + offset[:, 2] = torch.sin(phase * 2) * 0.1 # ±0.1 rad, double frequency + + # Add small random variation to make it more natural + noise = (torch.rand_like(init_pose, device=self.device) - 0.5) * 0.02 + + # Compute action + action = init_pose + offset + noise + + # Clamp to joint limits if available + if hasattr(self.robot.body_data, "qpos_limits"): + qpos_limits = self.robot.body_data.qpos_limits[0] # (num_joints, 2) + action = torch.clamp(action, qpos_limits[:, 0], qpos_limits[:, 1]) + + action_list.append(action) + + logger.log_info( + f"Generated {len(action_list)} demo actions with sinusoidal trajectory" + ) + return action_list diff --git a/embodichain_tasks/embodichain_tasks/tableware/__init__.py b/embodichain_tasks/embodichain_tasks/tableware/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/embodichain_tasks/embodichain_tasks/tableware/base_agent_env.py b/embodichain_tasks/embodichain_tasks/tableware/base_agent_env.py new file mode 100644 index 00000000..8814ea56 --- /dev/null +++ b/embodichain_tasks/embodichain_tasks/tableware/base_agent_env.py @@ -0,0 +1,201 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +import torch +from embodichain.utils import logger + + +class BaseAgentEnv: + def _init_agents(self, agent_config, task_name, agent_config_path=None): + from embodichain.agents.hierarchy.task_agent import TaskAgent + from embodichain.agents.hierarchy.code_agent import CodeAgent + from embodichain.agents.hierarchy.validation_agent import ValidationAgent + from embodichain.agents.hierarchy.llm import ( + task_llm, + code_llm, + validation_llm, + ) + + if agent_config.get("TaskAgent") is not None: + self.task_agent = TaskAgent( + task_llm, + **agent_config["Agent"], + **agent_config["TaskAgent"], + task_name=task_name, + config_dir=agent_config_path, + ) + self.code_agent = CodeAgent( + code_llm, + **agent_config["Agent"], + **agent_config.get("CodeAgent"), + task_name=task_name, + config_dir=agent_config_path, + ) + self.validation_agent = ValidationAgent( + validation_llm, + task_name=task_name, + task_description=self.code_agent.prompt_kwargs.get("task_prompt")[ + "content" + ], + basic_background=self.code_agent.prompt_kwargs.get("basic_background")[ + "content" + ], + atom_actions=self.code_agent.prompt_kwargs.get("atom_actions")["content"], + ) + + def get_states(self): + # TODO: only support num_env = 1 for now + # store robot states in each env.reset + self.init_qpos = self.robot.get_qpos().squeeze(0) + + self.left_arm_joints = self.robot.get_joint_ids(name="left_arm") + self.right_arm_joints = self.robot.get_joint_ids(name="right_arm") + self.left_eef_joints = self.robot.get_joint_ids(name="left_eef") + self.right_eef_joints = self.robot.get_joint_ids(name="right_eef") + + self.left_arm_init_qpos = self.init_qpos[self.left_arm_joints] + self.right_arm_init_qpos = self.init_qpos[self.right_arm_joints] + + self.left_arm_init_xpos = self.robot.compute_fk( + self.left_arm_init_qpos, name="left_arm", to_matrix=True + ).squeeze(0) + self.right_arm_init_xpos = self.robot.compute_fk( + self.right_arm_init_qpos, name="right_arm", to_matrix=True + ).squeeze(0) + + self.left_arm_current_qpos = self.left_arm_init_qpos + self.right_arm_current_qpos = self.right_arm_init_qpos + + self.left_arm_current_xpos = self.left_arm_init_xpos + self.right_arm_current_xpos = self.right_arm_init_xpos + + self.left_arm_base_pose = self.robot.get_control_part_base_pose( + "left_arm", to_matrix=True + ).squeeze(0) + self.right_arm_base_pose = self.robot.get_control_part_base_pose( + "right_arm", to_matrix=True + ).squeeze(0) + + self.open_state = torch.tensor([0.05]) + self.close_state = torch.tensor([0.0]) + self.left_arm_current_gripper_state = self.open_state + self.right_arm_current_gripper_state = self.open_state + + # store some useful obj information + init_obj_info = {} + obj_uids = self.sim.get_rigid_object_uid_list() + for obj_name in obj_uids: + obj = self.sim.get_rigid_object(obj_name) + obj_pose = obj.get_local_pose(to_matrix=True).squeeze(0) + obj_height = obj_pose[2, 3] # Extract the height (z-coordinate) + obj_grasp_pose = self.affordance_datas.get( + f"{obj_name}_grasp_pose_object", None + ) + init_obj_info[obj_name] = { + "pose": obj_pose, # Store the full pose (4x4 matrix) + "height": obj_height, # Store the height (z-coordinate) + "grasp_pose_obj": ( + obj_grasp_pose.squeeze(0) if obj_grasp_pose is not None else None + ), # Store the grasp pose if available + } + self.init_obj_info = init_obj_info + + # -------------------- Common getters / setters -------------------- + + def get_obs_for_agent(self): + obs = self.get_obs() + rgb = obs["sensor"]["cam_high"]["color"].squeeze(0) + + # Get validation camera data + camera_data = self.event_manager.get_functor("validation_cameras")(self, None) + result = {"rgb": rgb} + result.update({k: v.squeeze(0) for k, v in camera_data.items()}) + return result + + def get_current_qpos_agent(self): + return self.left_arm_current_qpos, self.right_arm_current_qpos + + def set_current_qpos_agent(self, arm_qpos, is_left): + if is_left: + self.left_arm_current_qpos = arm_qpos + else: + self.right_arm_current_qpos = arm_qpos + + def get_current_xpos_agent(self): + return self.left_arm_current_xpos, self.right_arm_current_xpos + + def set_current_xpos_agent(self, arm_xpos, is_left): + if is_left: + self.left_arm_current_xpos = arm_xpos + else: + self.right_arm_current_xpos = arm_xpos + + def get_current_gripper_state_agent(self): + return self.left_arm_current_gripper_state, self.right_arm_current_gripper_state + + def set_current_gripper_state_agent(self, arm_gripper_state, is_left): + if is_left: + self.left_arm_current_gripper_state = arm_gripper_state + else: + self.right_arm_current_gripper_state = arm_gripper_state + + # -------------------- IK / FK -------------------- + def get_arm_ik(self, target_xpos, is_left, qpos_seed=None): + control_part = "left_arm" if is_left else "right_arm" + ret, qpos = self.robot.compute_ik( + name=control_part, pose=target_xpos, joint_seed=qpos_seed + ) + return ret.all().item(), qpos.squeeze(0) + + def get_arm_fk(self, qpos, is_left): + control_part = "left_arm" if is_left else "right_arm" + xpos = self.robot.compute_fk( + name=control_part, qpos=torch.as_tensor(qpos), to_matrix=True + ) + return xpos.squeeze(0) + + # -------------------- get only code for action list -------------------- + def generate_code_for_actions(self, regenerate=False, **kwargs): + logger.log_info( + f"Generate code for creating action list for {self.code_agent.task_name}.", + color="green", + ) + + # Task planning + print(f"\033[92m\nStart task planning.\n\033[0m") + + task_agent_input = self.task_agent.get_composed_observations( + env=self, regenerate=regenerate, **kwargs + ) + task_plan = self.task_agent.generate(**task_agent_input) + + # Code generation + print(f"\033[94m\nStart code generation.\n\033[0m") + code_agent_input = self.code_agent.get_composed_observations( + env=self, regenerate=regenerate, **kwargs + ) + code_agent_input["task_plan"] = task_plan + + code_file_path, kwargs, code = self.code_agent.generate(**code_agent_input) + return code_file_path, kwargs, code + + # -------------------- get action list -------------------- + def create_demo_action_list(self, regenerate=False, *args, **kwargs): + code_file_path, kwargs, _ = self.generate_code_for_actions( + regenerate=regenerate + ) + action_list = self.code_agent.act(code_file_path, **kwargs) + return action_list diff --git a/embodichain_tasks/embodichain_tasks/tableware/blocks_ranking_rgb.py b/embodichain_tasks/embodichain_tasks/tableware/blocks_ranking_rgb.py new file mode 100644 index 00000000..b93ffd8d --- /dev/null +++ b/embodichain_tasks/embodichain_tasks/tableware/blocks_ranking_rgb.py @@ -0,0 +1,324 @@ +# ---------------------------------------------------------------------------- +# 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 +import numpy as np + +from embodichain.lab.gym.envs import EmbodiedEnv, EmbodiedEnvCfg +from embodichain.lab.gym.utils.registration import register_env +from embodichain.lab.sim.planners import ( + MotionGenerator, + MotionGenCfg, + MotionGenOptions, + ToppraPlannerCfg, + ToppraPlanOptions, + PlanState, + MoveType, + MovePart, +) +from embodichain.lab.sim.planners.utils import TrajectorySampleMethod +from embodichain.utils import logger + +__all__ = ["BlocksRankingRGBEnv"] + + +@register_env("BlocksRankingRGB-v1", max_episode_steps=600) +class BlocksRankingRGBEnv(EmbodiedEnv): + def __init__(self, cfg: EmbodiedEnvCfg = None, **kwargs): + super().__init__(cfg, **kwargs) + + action_config = kwargs.get("action_config", None) + if action_config is not None: + self.action_config = action_config + + def create_demo_action_list(self, *args, **kwargs): + """ + Create a demonstration action list for ranking three blocks in RGB order from left to right. + + Now the expert trajectory follows the following strategy: + - Do not move the green block block_2 (as the middle reference) + - Right hand (right_arm + right_eef) takes the red block block_1 and places it on the left of the green block + - Left hand (left_arm + left_eef) takes the blue block block_3 and places it on the right of the green block + + Returns: + list: A list of demo actions (torch.Tensor) to be executed by env.step(). + """ + try: + block1 = self.sim.get_rigid_object("block_1") # Red + block2 = self.sim.get_rigid_object("block_2") # Green + block3 = self.sim.get_rigid_object("block_3") # Blue + except Exception as e: + logger.log_warning(f"Blocks not found: {e}, returning empty action list.") + return [] + + # Get block poses and positions + b1_pose = block1.get_local_pose(to_matrix=True) + b2_pose = block2.get_local_pose(to_matrix=True) + b3_pose = block3.get_local_pose(to_matrix=True) + b1_pos = b1_pose[:, :3, 3] + b2_pos = b2_pose[:, :3, 3] + b3_pos = b3_pose[:, :3, 3] + + # Construct the target line centered on the green block (Red < Green < Blue) + base_x = b2_pos[:, 0] + base_y = b2_pos[:, 1] + base_z = b2_pos[:, 2] + + tgt_green = b2_pos # not moved, kept for reference + tgt_red = torch.stack([base_x - 0.12, base_y, base_z], dim=1) + tgt_blue = torch.stack([base_x + 0.10, base_y, base_z], dim=1) + + # Right arm / right hand + right_arm_ids = self.robot.get_joint_ids(name="right_arm") + right_eef_ids = self.robot.get_joint_ids(name="right_eef") + # Left arm / left hand + left_arm_ids = self.robot.get_joint_ids(name="left_arm") + left_eef_ids = self.robot.get_joint_ids(name="left_eef") + + init_qpos = self.robot.get_qpos() + init_right_arm_qpos = init_qpos[:, right_arm_ids] + init_right_arm_xpos = self.robot.compute_fk( + qpos=init_right_arm_qpos, name="right_arm", to_matrix=True + ) + init_left_arm_qpos = init_qpos[:, left_arm_ids] + init_left_arm_xpos = self.robot.compute_fk( + qpos=init_left_arm_qpos, name="left_arm", to_matrix=True + ) + + motion_cfg = MotionGenCfg( + planner_cfg=ToppraPlannerCfg( + robot_uid=self.robot.uid, + ) + ) + self.motion_generator = MotionGenerator(cfg=motion_cfg) + + gripper_open = torch.tensor( + [0.05, 0.05], dtype=torch.float32, device=self.device + ) + gripper_close = torch.tensor( + [0.0, 0.0], dtype=torch.float32, device=self.device + ) + + action_list = [] + + def _ik_to_qpos( + target_xpos: torch.Tensor, seed: torch.Tensor, arm_name: str, name: str + ): + is_success, qpos = self.robot.compute_ik( + pose=target_xpos, joint_seed=seed, name=arm_name + ) + success_flag = ( + is_success.all() if isinstance(is_success, torch.Tensor) else is_success + ) + if not success_flag: + logger.log_warning(f"IK failed for {name}, using previous qpos.") + qpos = seed + return qpos + + def _append_hold( + qpos: torch.Tensor, + num_steps: int, + gripper_state: torch.Tensor, + arm_ids, + eef_ids, + ): + for _ in range(num_steps): + action = init_qpos.clone() + action[:, arm_ids] = qpos + action[:, eef_ids] = gripper_state.unsqueeze(0).expand( + self.num_envs, -1 + ) + action_list.append(action) + + def _append_move_for_arm( + qpos_start: torch.Tensor, + qpos_end: torch.Tensor, + num_steps: int, + gripper_state: torch.Tensor, + arm_ids, + eef_ids, + control_part: str, + ): + options = MotionGenOptions( + control_part=control_part, + plan_opts=ToppraPlanOptions( + constraints={ + "velocity": 0.2, + "acceleration": 0.5, + }, + sample_method=TrajectorySampleMethod.QUANTITY, + sample_interval=num_steps, + ), + ) + # Joint space trajectory + target_states = [ + PlanState.single(qpos=qpos_start[0], move_type=MoveType.JOINT_MOVE), + PlanState.single( + qpos=qpos_end[0], + move_type=MoveType.JOINT_MOVE, + ), + ] + plan_result = self.motion_generator.generate( + target_states=target_states, options=options + ) + + for qpos_item in plan_result.positions[0]: + qpos = torch.as_tensor( + qpos_item, dtype=torch.float32, device=self.device + ) + qpos = qpos.flatten() + if qpos.shape[0] != len(arm_ids): + logger.log_warning( + f"Qpos shape mismatch: got {qpos.shape[0]}, expected {len(arm_ids)}" + ) + continue + qpos = qpos.unsqueeze(0).expand(self.num_envs, -1) + action = init_qpos.clone() + action[:, arm_ids] = qpos + action[:, eef_ids] = gripper_state.unsqueeze(0).expand( + self.num_envs, -1 + ) + action_list.append(action) + + def _pick_and_place( + block_pos: torch.Tensor, + place_pos: torch.Tensor, + seed_qpos: torch.Tensor, + init_arm_xpos: torch.Tensor, + arm_ids, + eef_ids, + arm_name: str, + tag: str, + ): + pick = init_arm_xpos.clone() + pick[:, :3, 3] = block_pos + torch.tensor( + [0.02, 0.0, -0.025], dtype=torch.float32, device=self.device + ).unsqueeze(0) + + lift = pick.clone() + lift[:, 2, 3] += 0.15 + + place = init_arm_xpos.clone() + place[:, :3, 3] = place_pos + torch.tensor( + [0.025, 0.0, 0.02], dtype=torch.float32, device=self.device + ).unsqueeze(0) + + # inverse kinematics for three key poses + q_pick = _ik_to_qpos(pick, seed_qpos, arm_name, f"{tag}_pick") + q_lift = _ik_to_qpos(lift, q_pick, arm_name, f"{tag}_lift") + q_place = _ik_to_qpos(place, q_lift, arm_name, f"{tag}_place") + + # execution segments: seed -> pick(open gripper) -> lift -> place(close gripper) + _append_move_for_arm( + seed_qpos, q_pick, 20, gripper_open, arm_ids, eef_ids, arm_name + ) + _append_hold(q_pick, 5, gripper_close, arm_ids, eef_ids) # close gripper + _append_move_for_arm( + q_pick, q_lift, 20, gripper_close, arm_ids, eef_ids, arm_name + ) + _append_move_for_arm( + q_lift, q_place, 30, gripper_close, arm_ids, eef_ids, arm_name + ) + _append_hold(q_place, 5, gripper_open, arm_ids, eef_ids) # open gripper + + return q_place + + # 1. Right hand handles the red block + current_seed_right = init_right_arm_qpos + current_seed_right = _pick_and_place( + b1_pos, + tgt_red, + current_seed_right, + init_right_arm_xpos, + right_arm_ids, + right_eef_ids, + "right_arm", + "red", + ) + + init_qpos[:, right_arm_ids] = current_seed_right + + # 2. Left hand handles the blue block + current_seed_left = init_left_arm_qpos + current_seed_left = _pick_and_place( + b3_pos, + tgt_blue, + current_seed_left, + init_left_arm_xpos, + left_arm_ids, + left_eef_ids, + "left_arm", + "blue", + ) + + logger.log_info(f"Generated {len(action_list)} demo actions for RGB ranking") + return action_list + + def is_task_success(self, **kwargs) -> torch.Tensor: + """Determine if the task is successfully completed. + + The task is successful if: + 1. Three blocks are arranged in RGB order from front to back: + - Red block (block_1) x < Green block (block_2) x < Blue block (block_3) x + 2. All blocks are close together (within tolerance) + + Args: + **kwargs: Additional arguments for task-specific success criteria. + + Returns: + torch.Tensor: A boolean tensor indicating success for each environment in the batch. + """ + try: + block1 = self.sim.get_rigid_object("block_1") # Red + block2 = self.sim.get_rigid_object("block_2") # Green + block3 = self.sim.get_rigid_object("block_3") # Blue + except Exception as e: + logger.log_warning(f"Blocks not found: {e}, returning False.") + return torch.zeros(self.num_envs, dtype=torch.bool, device=self.device) + + # Get block poses + block1_pose = block1.get_local_pose(to_matrix=True) + block2_pose = block2.get_local_pose(to_matrix=True) + block3_pose = block3.get_local_pose(to_matrix=True) + + # Extract positions (x, y, z) + block1_pos = block1_pose[:, :3, 3] # (num_envs, 3) + block2_pos = block2_pose[:, :3, 3] + block3_pos = block3_pose[:, :3, 3] + + # Tolerance for checking if blocks are close together + eps = torch.tensor([0.13, 0.03], dtype=torch.float32, device=self.device) + + # Check if blocks are close together in x-y plane + # block1 and block2 should be close + block1_block2_diff = torch.abs(block1_pos[:, :2] - block2_pos[:, :2]) + blocks_close_12 = torch.all(block1_block2_diff < eps.unsqueeze(0), dim=1) + + # block2 and block3 should be close + block2_block3_diff = torch.abs(block2_pos[:, :2] - block3_pos[:, :2]) + blocks_close_23 = torch.all(block2_block3_diff < eps.unsqueeze(0), dim=1) + + # Check RGB order: block1 (red) x < block2 (green) x < block3 (blue) x + rgb_order = (block1_pos[:, 0] < block2_pos[:, 0]) & ( + block2_pos[:, 0] < block3_pos[:, 0] + ) + + # Task succeeds if blocks are close together and in RGB order + success = blocks_close_12 & blocks_close_23 & rgb_order + + return success diff --git a/embodichain_tasks/embodichain_tasks/tableware/blocks_ranking_size.py b/embodichain_tasks/embodichain_tasks/tableware/blocks_ranking_size.py new file mode 100644 index 00000000..5d877c0f --- /dev/null +++ b/embodichain_tasks/embodichain_tasks/tableware/blocks_ranking_size.py @@ -0,0 +1,89 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +import torch +import numpy as np + +from embodichain.lab.gym.envs import EmbodiedEnv, EmbodiedEnvCfg +from embodichain.lab.gym.utils.registration import register_env +from embodichain.utils import logger + +__all__ = ["BlocksRankingSizeEnv"] + + +@register_env("BlocksRankingSize-v1", max_episode_steps=600) +class BlocksRankingSizeEnv(EmbodiedEnv): + def __init__(self, cfg: EmbodiedEnvCfg = None, **kwargs): + super().__init__(cfg, **kwargs) + + action_config = kwargs.get("action_config", None) + if action_config is not None: + self.action_config = action_config + + def is_task_success(self, **kwargs) -> torch.Tensor: + """Determine if the task is successfully completed. + + The task is successful if: + 1. Three blocks are arranged in size order from left to right: + - Large block (block_1) x < Medium block (block_2) x < Small block (block_3) x + 2. All blocks are close together (within tolerance) + + Args: + **kwargs: Additional arguments for task-specific success criteria. + + Returns: + torch.Tensor: A boolean tensor indicating success for each environment in the batch. + """ + try: + block1 = self.sim.get_rigid_object("block_1") # Large + block2 = self.sim.get_rigid_object("block_2") # Medium + block3 = self.sim.get_rigid_object("block_3") # Small + except Exception as e: + logger.log_warning(f"Blocks not found: {e}, returning False.") + return torch.zeros(self.num_envs, dtype=torch.bool, device=self.device) + + # Get block poses + block1_pose = block1.get_local_pose(to_matrix=True) + block2_pose = block2.get_local_pose(to_matrix=True) + block3_pose = block3.get_local_pose(to_matrix=True) + + # Extract positions (x, y, z) + block1_pos = block1_pose[:, :3, 3] # (num_envs, 3) + block2_pos = block2_pose[:, :3, 3] + block3_pos = block3_pose[:, :3, 3] + + # Tolerance for checking if blocks are close together + # Same as RoboTwin: eps = [0.13, 0.03] + eps = torch.tensor([0.13, 0.03], dtype=torch.float32, device=self.device) + + # Check if blocks are close together in x-y plane + # block1 and block2 should be close + block1_block2_diff = torch.abs(block1_pos[:, :2] - block2_pos[:, :2]) + blocks_close_12 = torch.all(block1_block2_diff < eps.unsqueeze(0), dim=1) + + # block2 and block3 should be close + block2_block3_diff = torch.abs(block2_pos[:, :2] - block3_pos[:, :2]) + blocks_close_23 = torch.all(block2_block3_diff < eps.unsqueeze(0), dim=1) + + # Check size order: block1 (large) x < block2 (medium) x < block3 (small) x + size_order = (block1_pos[:, 0] < block2_pos[:, 0]) & ( + block2_pos[:, 0] < block3_pos[:, 0] + ) + + # All conditions must be satisfied + success = blocks_close_12 & blocks_close_23 & size_order + + return success diff --git a/embodichain_tasks/embodichain_tasks/tableware/match_object_container.py b/embodichain_tasks/embodichain_tasks/tableware/match_object_container.py new file mode 100644 index 00000000..087a0611 --- /dev/null +++ b/embodichain_tasks/embodichain_tasks/tableware/match_object_container.py @@ -0,0 +1,176 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +import torch +import numpy as np + +from embodichain.lab.gym.envs import EmbodiedEnv, EmbodiedEnvCfg +from embodichain.lab.gym.utils.registration import register_env +from embodichain.utils import logger + +__all__ = ["MatchObjectContainerEnv"] + + +@register_env("MatchObjectContainer-v1", max_episode_steps=600) +class MatchObjectContainerEnv(EmbodiedEnv): + def __init__(self, cfg: EmbodiedEnvCfg = None, **kwargs): + super().__init__(cfg, **kwargs) + + action_config = kwargs.get("action_config", None) + if action_config is not None: + self.action_config = action_config + + def is_task_success(self, **kwargs) -> torch.Tensor: + """Determine if the task is successfully completed. + + This is a classification task: place blocks into matching shaped containers. + The task is successful if: + 1. Both cube blocks are inside container_cube (block_cube_1 and block_cube_2 -> container_cube) + 2. Both sphere blocks are inside container_sphere (block_sphere_1 and block_sphere_2 -> container_sphere) + 3. Both containers are up + + Args: + **kwargs: Additional arguments for task-specific success criteria. + + Returns: + torch.Tensor: A boolean tensor indicating success for each environment in the batch. + """ + try: + block_cube_1 = self.sim.get_rigid_object("block_cube_1") + block_sphere_1 = self.sim.get_rigid_object("block_sphere_1") + block_cube_2 = self.sim.get_rigid_object("block_cube_2") + block_sphere_2 = self.sim.get_rigid_object("block_sphere_2") + container_cube = self.sim.get_rigid_object("container_cube") + container_sphere = self.sim.get_rigid_object("container_sphere") + except Exception as e: + logger.log_warning(f"Blocks or containers not found: {e}, returning False.") + return torch.zeros(self.num_envs, dtype=torch.bool, device=self.device) + + # Get poses + block_cube_1_pose = block_cube_1.get_local_pose(to_matrix=True) + block_sphere_1_pose = block_sphere_1.get_local_pose(to_matrix=True) + block_cube_2_pose = block_cube_2.get_local_pose(to_matrix=True) + block_sphere_2_pose = block_sphere_2.get_local_pose(to_matrix=True) + container_cube_pose = container_cube.get_local_pose(to_matrix=True) + container_sphere_pose = container_sphere.get_local_pose(to_matrix=True) + + # Extract positions + block_cube_1_pos = block_cube_1_pose[:, :3, 3] # (num_envs, 3) + block_sphere_1_pos = block_sphere_1_pose[:, :3, 3] + block_cube_2_pos = block_cube_2_pose[:, :3, 3] + block_sphere_2_pos = block_sphere_2_pose[:, :3, 3] + container_cube_pos = container_cube_pose[:, :3, 3] + container_sphere_pos = container_sphere_pose[:, :3, 3] + + container_cube_fallen = self._is_fall(container_cube_pose) + container_sphere_fallen = self._is_fall(container_sphere_pose) + + # Check if blocks are inside their matching containers + # I got radius and height from the container_metal.obj file + container_bottom_radius = 0.1067 # Inner radius + z_tolerance = 0.05 # Vertical tolerance + container_height = 0.068 # Container height + container_half_height = container_height / 2 + + # Check if blocks are in their matching containers + cube_1_in_container = self._is_block_in_container( + block_cube_1_pos, + container_cube_pos, + container_bottom_radius, + container_half_height, + z_tolerance, + ) + cube_2_in_container = self._is_block_in_container( + block_cube_2_pos, + container_cube_pos, + container_bottom_radius, + container_half_height, + z_tolerance, + ) + sphere_1_in_container = self._is_block_in_container( + block_sphere_1_pos, + container_sphere_pos, + container_bottom_radius, + container_half_height, + z_tolerance, + ) + sphere_2_in_container = self._is_block_in_container( + block_sphere_2_pos, + container_sphere_pos, + container_bottom_radius, + container_half_height, + z_tolerance, + ) + + # Task success if cubes and spheres are in containers and containers are up + success = ( + cube_1_in_container + & cube_2_in_container + & sphere_1_in_container + & sphere_2_in_container + & ~container_cube_fallen + & ~container_sphere_fallen + ) + + return success + + def _is_block_in_container( + self, + block_pos: torch.Tensor, + container_pos: torch.Tensor, + container_bottom_radius: float, + container_half_height: float, + z_tolerance: float, + ) -> torch.Tensor: + """Check if a block is inside a container. + + Args: + block_pos: Block position (num_envs, 3) + container_pos: Container center position (num_envs, 3) + container_bottom_radius: Inner radius of container bottom in meters + container_half_height: Half height of container in meters + z_tolerance: Vertical tolerance for bottom check in meters + + Returns: + Boolean tensor indicating if block is in container (num_envs,) + """ + # XY plane distance check + xy_diff = torch.norm(block_pos[:, :2] - container_pos[:, :2], dim=1) + + # Z coordinate check: container center is at container_pos[:, 2] + # Container bottom = container_pos[:, 2] - container_half_height + # Container top = container_pos[:, 2] + container_half_height + z_above_bottom = ( + block_pos[:, 2] > container_pos[:, 2] - container_half_height - z_tolerance + ) + z_within_height = block_pos[:, 2] < container_pos[:, 2] + container_half_height + + return (xy_diff < container_bottom_radius) & z_above_bottom & z_within_height + + def _is_fall(self, pose: torch.Tensor) -> torch.Tensor: + # Extract z-axis from rotation matrix (last column, first 3 elements) + pose_rz = pose[:, :3, 2] + world_z_axis = torch.tensor([0, 0, 1], dtype=pose.dtype, device=pose.device) + + # Compute dot product for each batch element + dot_product = torch.sum(pose_rz * world_z_axis, dim=-1) # Shape: (batch_size,) + + # Clamp to avoid numerical issues with arccos + dot_product = torch.clamp(dot_product, -1.0, 1.0) + + # Compute angle and check if fallen + angle = torch.arccos(dot_product) + return angle >= torch.pi / 4 diff --git a/embodichain_tasks/embodichain_tasks/tableware/place_object_drawer.py b/embodichain_tasks/embodichain_tasks/tableware/place_object_drawer.py new file mode 100644 index 00000000..1d7da3b0 --- /dev/null +++ b/embodichain_tasks/embodichain_tasks/tableware/place_object_drawer.py @@ -0,0 +1,83 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +import torch +import numpy as np + + +from embodichain.lab.gym.envs import EmbodiedEnv, EmbodiedEnvCfg +from embodichain.lab.gym.utils.registration import register_env +from embodichain.utils import logger + +__all__ = ["PlaceObjectDrawerEnv"] + + +@register_env("PlaceObjectDrawer-v1", max_episode_steps=600) +class PlaceObjectDrawerEnv(EmbodiedEnv): + def __init__(self, cfg: EmbodiedEnvCfg = None, **kwargs): + super().__init__(cfg, **kwargs) + + action_config = kwargs.get("action_config", None) + if action_config is not None: + self.action_config = action_config + + def is_task_success(self, **kwargs) -> torch.Tensor: + """Determine if the task is successfully completed. + + The task is successful if: + 1. Object is within drawer inner_box area + 2. Drawer has been closed + + Args: + **kwargs: Additional arguments for task-specific success criteria. + + Returns: + torch.Tensor: A boolean tensor indicating success for each environment in the batch. + """ + try: + object_obj = self.sim.get_rigid_object("object") + drawer = self.sim.get_articulation("drawer") + except Exception as e: + logger.log_warning(f"Object or drawer not found: {e}, returning False.") + return torch.zeros(self.num_envs, dtype=torch.bool, device=self.device) + + # Get poses + object_pose = object_obj.get_local_pose(to_matrix=True) + # Get drawer inner_box (drawer interior) pose, not outer_box + inner_box_pose = drawer.get_link_pose("inner_box", to_matrix=True) + + # Extract positions + object_pos = object_pose[:, :3, 3] # (num_envs, 3) + inner_box_pos = inner_box_pose[:, :3, 3] # (num_envs, 3) + + # Get drawer joint position + drawer_qpos = drawer.get_qpos() # (num_envs, num_joints) + drawer_joint_pos = drawer_qpos[:, 0] + + # Check if drawer has been closed + drawer_closed = drawer_joint_pos < 0.05 + + # Check if object is within drawer inner_box area + xy_diff = torch.abs(object_pos[:, :2] - inner_box_pos[:, :2]) # (num_envs, 2) + xy_tolerance = torch.tensor( + [0.03, 0.03], dtype=torch.float32, device=self.device + ) + object_in_drawer_xy = torch.all(xy_diff < xy_tolerance.unsqueeze(0), dim=1) + + # Object must be in drawer and drawer must be closed + success = drawer_closed & object_in_drawer_xy + + return success diff --git a/embodichain_tasks/embodichain_tasks/tableware/pour_water/__init__.py b/embodichain_tasks/embodichain_tasks/tableware/pour_water/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/embodichain_tasks/embodichain_tasks/tableware/pour_water/action_bank.py b/embodichain_tasks/embodichain_tasks/tableware/pour_water/action_bank.py new file mode 100644 index 00000000..e68fbf78 --- /dev/null +++ b/embodichain_tasks/embodichain_tasks/tableware/pour_water/action_bank.py @@ -0,0 +1,255 @@ +# ---------------------------------------------------------------------------- +# 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 +import numpy as np +from copy import deepcopy +from typing import Dict, List +from embodichain.lab.gym.envs.action_bank.configurable_action import ( + ActionBank, + tag_node, + tag_edge, +) + +from embodichain.lab.gym.utils.misc import ( + resolve_env_params, + mul_linear_expand, + get_offset_pose_list, + get_changed_pose, +) + +from embodichain.lab.sim.planners import ( + MoveType, + PlanState, + MotionGenerator, + MotionGenCfg, + MotionGenOptions, + ToppraPlanOptions, + ToppraPlannerCfg, +) +from embodichain.utils import logger + +__all__ = ["PourWaterActionBank"] + + +class PourWaterActionBank(ActionBank): + @staticmethod + @tag_node + @resolve_env_params + def generate_left_arm_aim_qpos( + env, + valid_funcs_name_kwargs_proc: List | None = None, + ): + # FIXME FIXME FIXME FIXME + logger.log_warning( + f"CAUTION=============================THIS FUNC generate_left_arm_aim_qpos IS WRONG!!!! PLEASE FIX IT!!!!" + ) + left_aim_horizontal_angle = np.arctan2( + *( + ( + env.affordance_datas["cup_pose"][:2, 3] + - env.affordance_datas["left_arm_base_pose"][:2, 3] + )[1::-1] + ) + ) + left_arm_aim_qpos = deepcopy(env.affordance_datas["left_arm_init_qpos"]) + left_arm_aim_qpos[0] = left_aim_horizontal_angle + env.affordance_datas["left_arm_aim_qpos"] = left_arm_aim_qpos + return True + + @staticmethod + @tag_node + @resolve_env_params + # DONE: valid & process qpos & fk + def generate_right_arm_aim_qpos( + env, + valid_funcs_name_kwargs_proc: list | None = None, + ): + # FIXME FIXME FIXME FIXME + logger.log_warning( + f"CAUTION=============================THIS FUNC generate_right_arm_aim_qpos IS WRONG!!!! PLEASE FIX IT!!!!" + ) + right_aim_horizontal_angle = np.arctan2( + *( + ( + env.affordance_datas["bottle_pose"][:2, 3] + - env.affordance_datas["right_arm_base_pose"][:2, 3] + )[1::-1] + ) + ) + right_arm_aim_qpos = deepcopy(env.affordance_datas["right_arm_init_qpos"]) + right_arm_aim_qpos[0] = right_aim_horizontal_angle + env.affordance_datas["right_arm_aim_qpos"] = right_arm_aim_qpos + return True + + @staticmethod + @tag_node + @resolve_env_params + def compute_unoffset_for_exp(env, pose_input_output_names_changes: Dict = {}): + env.affordance_datas["bottle_grasp_unoffset_matrix_object"] = np.eye( + 4 + ) # For the overall transform matrix calculation + for input_pose_name, change_params in pose_input_output_names_changes.items(): + output_pose_name = change_params["output_pose_name"] + pose_changes = change_params["pose_changes"] + env.affordance_datas[output_pose_name] = get_changed_pose( + env.affordance_datas[input_pose_name], pose_changes + ) + + return True + + @staticmethod + @tag_edge + @tag_node + # TODO: Got the dimension from the scope + def execute_open(env, return_action: bool = False, **kwargs): + if return_action: + duration = kwargs.get("duration", 1) + expand = kwargs.get("expand", False) + if expand: + action = mul_linear_expand(np.array([[0.0], [1.0]]), [duration - 1]) + action = np.concatenate([action, np.array([[1.0]])]).transpose() + else: + action = np.ones((1, duration)) + return action + else: + return True + + @staticmethod + @tag_edge + @tag_node + def execute_close(env, return_action: bool = False, **kwargs): + + if return_action: + duration = kwargs.get("duration", 1) + expand = kwargs.get("expand", False) + if expand: + action = mul_linear_expand(np.array([[1.0], [0.0]]), [duration - 1]) + action = np.concatenate([action, np.array([[0.0]])]).transpose() + else: + action = np.zeros((1, duration)) + return action + else: + return True + + @staticmethod + @tag_edge + def plan_trajectory( + env, + agent_uid: str, + keypose_names: List[str], + duration: int, + edge_name: str = "", + ): + keyposes = [ + env.affordance_datas[keypose_name] for keypose_name in keypose_names + ] + + keyposes = [ + kp.cpu().numpy() if hasattr(kp, "cpu") and hasattr(kp, "numpy") else kp + for kp in keyposes + ] + + if all( + np.linalg.norm(former - latter).sum() <= 1e-3 + for former, latter in zip(keyposes, keyposes[1:]) + ): + logger.log_warning( + f"Applying plan_trajectory to two very close qpos! Using stand_still." + ) + keyposes = [keyposes[0]] * 2 + ret_transposed = PourWaterActionBank.stand_still( + env, + agent_uid, + keypose_names, + duration, + ) + + return ret_transposed + + else: + motion_generator = MotionGenerator( + cfg=MotionGenCfg(planner_cfg=ToppraPlannerCfg(robot_uid=env.robot.uid)) + ) + + plan_state = [ + PlanState.single( + qpos=torch.as_tensor(qpos), move_type=MoveType.JOINT_MOVE + ) + for qpos in keyposes + ] + + ret = motion_generator.generate( + target_states=plan_state, + options=MotionGenOptions( + control_part=agent_uid, + plan_opts=ToppraPlanOptions( + sample_interval=duration, + ), + ), + ) + + return ret.positions[0].detach().cpu().numpy().T + + @staticmethod + @tag_edge + def stand_still( + env, + agent_uid: str, + keypose_names: List[str], + duration: int, + ): + keyposes = [ + env.affordance_datas[keypose_name] for keypose_name in keypose_names + ] + + stand_still_qpos = keyposes[0] + + if ( + stand_still_qpos.shape + != np.asarray(env.robot.get_joint_ids("left_arm")).shape + ): + logger.log_error( + f"The shape of stand_still qpos is different from {agent_uid}'s setting." + ) + + if any( + np.linalg.norm(former - latter).sum() > 1e-6 + for former, latter in zip(keyposes, keyposes[1:]) + ): + logger.log_warning( + f"Applying stand still to two different qpos! Using the first qpos {stand_still_qpos}" + ) + keyposes = [stand_still_qpos] * 2 + + ret = np.asarray([stand_still_qpos] * duration) + + return ret.T + + @staticmethod + @tag_edge + def left_arm_go_back(env, duration: int): + left_arm_monitor_qpos, left_arm_init_qpos = ( + env.affordance_datas["left_arm_monitor_qpos"], + env.affordance_datas["left_arm_init_qpos"], + ) + left_home_sample_num = duration + qpos_expand_left = np.array([left_arm_monitor_qpos, left_arm_init_qpos]) + qpos_expand_left = mul_linear_expand(qpos_expand_left, [left_home_sample_num]) + ret = np.array(qpos_expand_left).T + return ret diff --git a/embodichain_tasks/embodichain_tasks/tableware/pour_water/pour_water.py b/embodichain_tasks/embodichain_tasks/tableware/pour_water/pour_water.py new file mode 100644 index 00000000..cff455d2 --- /dev/null +++ b/embodichain_tasks/embodichain_tasks/tableware/pour_water/pour_water.py @@ -0,0 +1,161 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +import torch +from typing import Dict, Optional + +from embodichain.lab.gym.envs import EmbodiedEnv, EmbodiedEnvCfg +from embodichain.lab.gym.utils.registration import register_env +from embodichain.utils import logger + +from embodichain_tasks.tableware.base_agent_env import BaseAgentEnv +from embodichain_tasks.tableware.pour_water.action_bank import ( + PourWaterActionBank, +) + +__all__ = ["PourWaterEnv", "PourWaterAgentEnv"] + + +@register_env("PourWater-v3", max_episode_steps=600) +class PourWaterEnv(EmbodiedEnv): + def __init__(self, cfg: EmbodiedEnvCfg = None, **kwargs): + super().__init__(cfg, **kwargs) + + action_config = kwargs.get("action_config", None) + if action_config is not None: + self.action_config = action_config + + def create_demo_action_list(self, *args, **kwargs): + """ + Create a demonstration action list for the current task. + + Returns: + list: A list of demo actions generated by the task. + """ + logger.log_info("Create demo action list for PourWaterTask.") + + if getattr(self, "action_config") is not None: + self._init_action_bank(PourWaterActionBank, self.action_config) + action_list = self.create_expert_demo_action_list(*args, **kwargs) + else: + logger.log_error("No action_config found in env, please check again.") + + if action_list is None: + return action_list + + logger.log_info( + f"Demo action list created with {len(action_list)} steps.", color="green" + ) + return action_list + + def create_expert_demo_action_list(self, **kwargs): + """ + Create an expert demonstration action list using the action bank. + + This function generates a trajectory based on expert knowledge, mapping joint and end-effector + states to the required action format for the environment and robot type. + + Args: + **kwargs: Additional keyword arguments. + + Returns: + list: A list of actions, each containing joint positions ("qpos"). + """ + + if hasattr(self, "action_bank") is False or self.action_bank is None: + logger.log_error( + "Action bank is not initialized. Cannot create expert demo action list." + ) + + ret = self.action_bank.create_action_list( + self, self.graph_compose, self.packages + ) + + if ret is None: + logger.log_warning("Failed to generate expert demo action list.") + return None + + # TODO: to be removed, need a unified interface in robot class + left_arm_joints = self.robot.get_joint_ids(name="left_arm") + right_arm_joints = self.robot.get_joint_ids(name="right_arm") + left_eef_joints = self.robot.get_joint_ids(name="left_eef", remove_mimic=True) + right_eef_joints = self.robot.get_joint_ids(name="right_eef", remove_mimic=True) + + total_traj_num = ret[list(ret.keys())[0]].shape[-1] + actions = torch.zeros( + (total_traj_num, self.num_envs, self.robot.dof), dtype=torch.float32 + ) + + for key, joints in [ + ("left_arm", left_arm_joints), + ("left_eef", left_eef_joints), + ("right_arm", right_arm_joints), + ("right_eef", right_eef_joints), + ]: + if key in ret: + # TODO: only 1 env supported now + actions[:, 0, joints] = torch.as_tensor(ret[key].T, dtype=torch.float32) + + return actions + + def is_task_success(self, **kwargs) -> torch.Tensor: + """Determine if the task is successfully completed. This is mainly used in the data generation process + of the imitation learning. + + Args: + **kwargs: Additional arguments for task-specific success criteria. + + Returns: + torch.Tensor: A boolean tensor indicating success for each environment in the batch. + """ + + bottle = self.sim.get_rigid_object("bottle") + cup = self.sim.get_rigid_object("cup") + + bottle_final_xpos = bottle.get_local_pose(to_matrix=True) + cup_final_xpos = cup.get_local_pose(to_matrix=True) + + bottle_ret = self._is_fall(bottle_final_xpos) + cup_ret = self._is_fall(cup_final_xpos) + + return ~(bottle_ret | cup_ret) + + def _is_fall(self, pose: torch.Tensor) -> torch.Tensor: + # Extract z-axis from rotation matrix (last column, first 3 elements) + pose_rz = pose[:, :3, 2] + world_z_axis = torch.tensor([0, 0, 1], dtype=pose.dtype, device=pose.device) + + # Compute dot product for each batch element + dot_product = torch.sum(pose_rz * world_z_axis, dim=-1) # Shape: (batch_size,) + + # Clamp to avoid numerical issues with arccos + dot_product = torch.clamp(dot_product, -1.0, 1.0) + + # Compute angle and check if fallen + angle = torch.arccos(dot_product) + return angle >= torch.pi / 4 + + +@register_env("PourWaterAgent-v3", max_episode_steps=600) +class PourWaterAgentEnv(BaseAgentEnv, PourWaterEnv): + def __init__(self, cfg: EmbodiedEnvCfg = None, **kwargs): + super().__init__(cfg, **kwargs) + super()._init_agents(**kwargs) + + def reset(self, seed: Optional[int] = None, options: Optional[Dict] = None): + obs, info = super().reset(seed=seed, options=options) + super().get_states() + return obs, info diff --git a/embodichain_tasks/embodichain_tasks/tableware/rearrangement.py b/embodichain_tasks/embodichain_tasks/tableware/rearrangement.py new file mode 100644 index 00000000..0658ab3a --- /dev/null +++ b/embodichain_tasks/embodichain_tasks/tableware/rearrangement.py @@ -0,0 +1,90 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# All rights reserved. +# ---------------------------------------------------------------------------- + +from typing import Dict, Optional +from embodichain.lab.gym.envs import EmbodiedEnv, EmbodiedEnvCfg +from embodichain.lab.gym.utils.registration import register_env +from embodichain_tasks.tableware.base_agent_env import BaseAgentEnv + +__all__ = ["RearrangementEnv", "RearrangementAgentEnv"] + + +@register_env("Rearrangement-v3", max_episode_steps=600) +class RearrangementEnv(EmbodiedEnv): + def __init__(self, cfg: EmbodiedEnvCfg = None, **kwargs): + super().__init__(cfg, **kwargs) + + action_config = kwargs.get("action_config", None) + if action_config is not None: + self.action_config = action_config + + def is_task_success(self) -> bool: + fork = self.sim.get_rigid_object("fork") + spoon = self.sim.get_rigid_object("spoon") + plate = self.sim.get_rigid_object("plate") + plate_pose = plate.get_local_pose(to_matrix=True) + # TODO: now only for 1 env + ( + spoon_place_target_x, + spoon_place_target_y, + spoon_place_target_z, + ) = self.affordance_datas["spoon_place_pose"][:3, 3] + ( + fork_place_target_x, + fork_place_target_y, + fork_place_target_z, + ) = self.affordance_datas["fork_place_pose"][:3, 3] + + spoon_pose = spoon.get_local_pose(to_matrix=True) + spoon_x, spoon_y, spoon_z = spoon_pose[0, :3, 3] + + fork_pose = fork.get_local_pose(to_matrix=True) + fork_x, fork_y, fork_z = fork_pose[0, :3, 3] + + tolerance = self.metadata.get("success_params", {}).get("tolerance", 0.02) + + # spoon and fork should with the x y range of tolerance related to plate. + return ~( + abs(spoon_x - spoon_place_target_x) > tolerance + or abs(spoon_y - spoon_place_target_y) > tolerance + or abs(fork_x - fork_place_target_x) > tolerance + or abs(fork_y - fork_place_target_y) > tolerance + ) + + +@register_env("RearrangementAgent-v3", max_episode_steps=600) +class RearrangementAgentEnv(BaseAgentEnv, RearrangementEnv): + def __init__(self, cfg: EmbodiedEnvCfg = None, **kwargs): + super().__init__(cfg, **kwargs) + super()._init_agents(**kwargs) + + def reset(self, seed: Optional[int] = None, options: Optional[Dict] = None): + obs, info = super().reset(seed=seed, options=options) + super().get_states() + return obs, info + + def is_task_success(self): + fork = self.sim.get_rigid_object("fork") + spoon = self.sim.get_rigid_object("spoon") + plate = self.sim.get_rigid_object("plate") + + plate_pose = plate.get_local_pose(to_matrix=True) + spoon_place_target_y = plate_pose[0, 1, 3] - 0.16 + fork_place_target_y = plate_pose[0, 1, 3] + 0.16 + + spoon_pose = spoon.get_local_pose(to_matrix=True) + spoon_y = spoon_pose[0, 1, 3] + + fork_pose = fork.get_local_pose(to_matrix=True) + fork_y = fork_pose[0, 1, 3] + + tolerance = self.metadata.get("success_params", {}).get("tolerance", 0.02) + + # spoon and fork should with the y range of tolerance related to plate. + return ( + abs(spoon_y - spoon_place_target_y) <= tolerance + and abs(fork_y - fork_place_target_y) <= tolerance + ) diff --git a/embodichain_tasks/embodichain_tasks/tableware/scoop_ice.py b/embodichain_tasks/embodichain_tasks/tableware/scoop_ice.py new file mode 100644 index 00000000..2da7536a --- /dev/null +++ b/embodichain_tasks/embodichain_tasks/tableware/scoop_ice.py @@ -0,0 +1,213 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +import os +import torch +import numpy as np +import pickle + +from copy import deepcopy +from typing import Sequence +from scipy.spatial.transform import Rotation as R + +from embodichain.lab.gym.envs import EmbodiedEnv, EmbodiedEnvCfg +from embodichain.lab.gym.utils.registration import register_env +from embodichain.data import get_data_path +from embodichain.utils import logger +from tqdm import tqdm + + +@register_env("ScoopIce-v1", max_episode_steps=600) +class ScoopIce(EmbodiedEnv): + def __init__(self, cfg: EmbodiedEnvCfg = None, **kwargs): + super().__init__(cfg, **kwargs) + + self.affordance_datas = {} + + # TODO: hardcode code, should be implemented as functor way. + self.trajectory = pickle.load( + open( + get_data_path("ScoopIceNewEnv/pose_record_20250919_184544.pkl"), + "rb", + ) + ) + self.trajectory_sample_rate = 2 + + def set_scoop_pose(self, xyzrxryrz): + scoop = self.sim.get_rigid_object("scoop") + pose = np.eye(4) + pose[:3, 3] = xyzrxryrz[:3] + pose[:3, :3] = R.from_euler("XYZ", xyzrxryrz[3:], degrees=True).as_matrix() + n_env = self.sim.num_envs + pose_t = torch.tensor( + pose[None, :, :].repeat(n_env, axis=0), + dtype=torch.float32, + device=self.device, + ) + scoop.set_local_pose(pose_t) + + def set_cup_pose(self, xyzrxryrz): + cup = self.sim.get_rigid_object("paper_cup") + pose = np.eye(4) + pose[:3, 3] = xyzrxryrz[:3] + pose[:3, :3] = R.from_euler("XYZ", xyzrxryrz[3:], degrees=True).as_matrix() + n_env = self.sim.num_envs + pose_t = torch.tensor( + pose[None, :, :].repeat(n_env, axis=0), + dtype=torch.float32, + device=self.device, + ) + cup.set_local_pose(pose_t) + + def add_xpos_offset(self, arm_qpos: np.ndarray, offset: np.ndarray, is_left: bool): + """Add offset to arm qposes along end-effector x axis. + + Args: + arm_qposes (np.ndarray): [waypoint_num, dof] + """ + waypoint_num = arm_qpos.shape[0] + dof = arm_qpos.shape[1] + offset_t = torch.tensor(offset, dtype=torch.float32, device=self.device) + control_part = "left_arm" if is_left else "right_arm" + + arm_qpos_batch = torch.tensor( + arm_qpos[None, :, :], dtype=torch.float32, device=self.device + ) + + arm_xpos_batch = self.robot.compute_batch_fk( + qpos=arm_qpos_batch, name=control_part, to_matrix=True + ) + arm_xpos_batch[:, :, :3, 3] += offset_t + ret, arm_qpos_offset_batch = self.robot.compute_batch_ik( + pose=arm_xpos_batch, + joint_seed=arm_qpos_batch, + name=control_part, + ) + return arm_qpos_offset_batch[0].to("cpu").numpy() + + def pack_qpos(self): + left_arm_qpos = self.trajectory["left_arm"] # [waypoint_num, dof] + logger.log_info("Adding x and z offset to left arm trajectory...") + left_arm_qpos = self.add_xpos_offset( + arm_qpos=left_arm_qpos, offset=np.array([-0.018, 0.0, -0.01]), is_left=True + ) + right_arm_qpos = self.trajectory["right_arm"] # [waypoint_num, dof] + # TODO: add z offset to right arm + logger.log_info("Adding z offset to right arm trajectory...") + right_arm_qpos = self.add_xpos_offset( + arm_qpos=right_arm_qpos, offset=np.array([0.00, 0.0, 0.02]), is_left=False + ) + left_eef_qpos = self.trajectory["left_eef"] # [waypoint_num, hand_dof] + right_eef_qpos = self.trajectory["right_eef"] + torso_qpos = self.trajectory["torso"] + # TODO: need head qpos. + + left_arm_qpos_expand = left_arm_qpos[None, :, :].repeat(self.num_envs, axis=0) + right_arm_qpos_expand = right_arm_qpos[None, :, :].repeat(self.num_envs, axis=0) + left_eef_qpos_expand = left_eef_qpos[None, :, :].repeat(self.num_envs, axis=0) + right_eef_qpos_expand = right_eef_qpos[None, :, :].repeat(self.num_envs, axis=0) + torso_qpos_expand = torso_qpos[None, :, :].repeat(self.num_envs, axis=0) + all_qpos = np.concatenate( + [ + left_arm_qpos_expand, + right_arm_qpos_expand, + left_eef_qpos_expand, + right_eef_qpos_expand, + torso_qpos_expand, + ], + axis=2, + ) + return all_qpos + + def _initialize_episode( + self, env_ids: Sequence[int] | None = None, **kwargs + ) -> None: + + left_arm_ids = self.robot.get_joint_ids(name="left_arm") + right_arm_ids = self.robot.get_joint_ids(name="right_arm") + left_eef_ids = self.robot.get_joint_ids(name="left_eef") + right_eef_ids = self.robot.get_joint_ids(name="right_eef") + torso_ids = self.robot.get_joint_ids(name="torso") + all_ids = np.hstack( + [left_arm_ids, right_arm_ids, left_eef_ids, right_eef_ids, torso_ids] + ) + + # TODO: read xy random range from config + xy_random_range = np.array([[-0.01, -0.01], [0.01, 0.01]]) + xy_random_offset = np.zeros(shape=(self.num_envs, 2)) + for arena_id in range(self.num_envs): + xy_random_offset[arena_id] = np.random.uniform( + low=xy_random_range[0], high=xy_random_range[1], size=(2,) + ) + # TODO: apply warping to container pose + + all_qpos = self.pack_qpos() + all_qpos_t = torch.tensor(all_qpos, dtype=torch.float32, device=self.device) + + # to initial qpos + left_open_qpos = np.array([0.06, 1.5, 0.2, 0.2, 0.2, 0.2]) + left_close_qpos = np.array([0.13, 1.5, 0.5, 0.5, 0.5, 0.5]) + right_open_qpos = np.array([0.3, 1.5, 0.3, 0.3, 0.3, 0.3]) + right_close_qpos = np.array([0.6, 1.5, 0.7, 0.5, 0.7, 0.6]) + + all_qpos_t[:, :, 14:20] = torch.tensor( + left_close_qpos, dtype=torch.float32, device=self.device + ) + all_qpos_t[:, :, 20:26] = torch.tensor( + right_close_qpos, dtype=torch.float32, device=self.device + ) + + first_close_qpos = all_qpos_t[:, 0, :].to("cpu").numpy() + first_open_qpos = deepcopy(first_close_qpos) + + # to first open pose + first_open_qpos[:, 14:20] = left_open_qpos + first_open_qpos[:, 20:26] = right_open_qpos + self.robot.set_qpos( + torch.tensor(first_open_qpos, dtype=torch.float32, device=self.device), + joint_ids=all_ids, + ) + self.sim.update(step=200) + # save warp trajectory as demo action list + waypoint_num = self.trajectory["left_arm"].shape[0] + current_qpos = self.robot.get_qpos() + self.demo_action_list = [] + for waypoint_idx in range(waypoint_num): + action = current_qpos.clone() + action[:, all_ids] = all_qpos_t[:, waypoint_idx, :] + # TODO: sample in trajectory + self.demo_action_list.append(action) + + # TODO: tricky implementation. Hold the first joint state for a while. + if waypoint_idx == 0: + for _ in range(20): + self.demo_action_list.append(action) + + self.sim.update(step=100) + + # apply events such as randomization for environments that need a reset + if self.cfg.events: + if "reset" in self.event_manager.available_modes: + self.event_manager.apply(mode="reset", env_ids=env_ids) + + def create_demo_action_list(self, *args, **kwargs): + logger.log_info( + f"The original demo action list length: {len(self.demo_action_list)}" + ) + logger.log_info( + f"Downsample the demo action list by self.trajectory_sample_rate5 times." + ) + return self.demo_action_list[:: self.trajectory_sample_rate] diff --git a/embodichain_tasks/embodichain_tasks/tableware/stack_blocks_two.py b/embodichain_tasks/embodichain_tasks/tableware/stack_blocks_two.py new file mode 100644 index 00000000..ca1fc1fb --- /dev/null +++ b/embodichain_tasks/embodichain_tasks/tableware/stack_blocks_two.py @@ -0,0 +1,324 @@ +# ---------------------------------------------------------------------------- +# 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, express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +from __future__ import annotations + +import torch +import numpy as np + +from embodichain.lab.gym.envs import EmbodiedEnv, EmbodiedEnvCfg +from embodichain.lab.gym.utils.registration import register_env +from embodichain.lab.sim.planners import ( + MotionGenerator, + MotionGenCfg, + MotionGenOptions, + ToppraPlannerCfg, + ToppraPlanOptions, + PlanState, + MoveType, + MovePart, +) +from embodichain.lab.sim.planners.utils import TrajectorySampleMethod +from embodichain.utils import logger + +__all__ = ["StackBlocksTwoEnv"] + + +@register_env("StackBlocksTwo-v1", max_episode_steps=600) +class StackBlocksTwoEnv(EmbodiedEnv): + def __init__(self, cfg: EmbodiedEnvCfg = None, **kwargs): + super().__init__(cfg, **kwargs) + + action_config = kwargs.get("action_config", None) + if action_config is not None: + self.action_config = action_config + + def create_demo_action_list(self, *args, **kwargs): + """ + Create a demonstration action list for stacking block_2 on top of block_1. + + This expert trajectory: + 1. Moves to block_2 position + 2. Grasps block_2 + 3. Lifts block_2 + 4. Moves to block_1 position + 5. Places block_2 on top of block_1 + 6. Releases and retracts + + Returns: + list: A list of demo actions generated by the task. + """ + try: + block1 = self.sim.get_rigid_object("block_1") + block2 = self.sim.get_rigid_object("block_2") + except Exception as e: + logger.log_warning(f"Blocks not found: {e}, returning empty action list.") + return [] + + # Get block positions + block1_pose = block1.get_local_pose(to_matrix=True) + block2_pose = block2.get_local_pose(to_matrix=True) + + block1_pos = block1_pose[:, :3, 3] # (num_envs, 3) + block2_pos = block2_pose[:, :3, 3] # (num_envs, 3) + + right_arm_ids = self.robot.get_joint_ids(name="right_arm") + right_eef_ids = self.robot.get_joint_ids(name="right_eef") + + init_qpos = self.robot.get_qpos() + init_right_arm_qpos = init_qpos[:, right_arm_ids] + + init_right_arm_xpos = self.robot.compute_fk( + qpos=init_right_arm_qpos, name="right_arm", to_matrix=True + ) + + # 1. Approach block_2 + approach_block2_xpos = init_right_arm_xpos.clone() + approach_block2_xpos[:, :3, 3] = block2_pos + torch.tensor( + [0.02, 0.0, 0.1], dtype=torch.float32, device=self.device + ).unsqueeze(0) + + # 2. Pick block_2 + pick_block2_xpos = init_right_arm_xpos.clone() + pick_block2_xpos[:, :3, 3] = block2_pos + torch.tensor( + [0.02, 0.0, -0.025], dtype=torch.float32, device=self.device + ).unsqueeze(0) + + # 3. Lift block_2 + lift_block2_xpos = pick_block2_xpos.clone() + lift_block2_xpos[:, 2, 3] += 0.15 + + # 4. Approach block_1 + approach_block1_xpos = init_right_arm_xpos.clone() + target_place_pos = block1_pos.clone() + target_place_pos[:, 2] += 0.05 # block_1 height + block_2 height + approach_block1_xpos[:, :3, 3] = target_place_pos + torch.tensor( + [0.025, 0.0, 0.05], dtype=torch.float32, device=self.device + ).unsqueeze(0) + + # 5. Place block_2 on block_1 + place_block2_xpos = approach_block1_xpos.clone() + place_block2_xpos[:, :3, 3] = target_place_pos + torch.tensor( + [0.025, 0.0, 0.02], dtype=torch.float32, device=self.device + ).unsqueeze(0) + + # 6. Retract + retract_xpos = approach_block1_xpos.clone() + + retract_xpos[:, 2, 3] += 0.1 + + # Compute IK for each waypoint + waypoints = [ + ("approach_block2", approach_block2_xpos), + ("pick_block2", pick_block2_xpos), + ("lift_block2", lift_block2_xpos), + ("approach_block1", approach_block1_xpos), + ("place_block2", place_block2_xpos), + ("retract", retract_xpos), + ] + + qpos_waypoints = [] + current_qpos = init_right_arm_qpos + + for waypoint_name, waypoint_xpos in waypoints: + is_success, qpos = self.robot.compute_ik( + pose=waypoint_xpos, joint_seed=current_qpos, name="right_arm" + ) + + # Check IK success + success_flag = ( + is_success.all() if isinstance(is_success, torch.Tensor) else is_success + ) + if not success_flag: + logger.log_warning(f"IK failed for {waypoint_name}") + qpos = current_qpos + + qpos_waypoints.append(qpos) + current_qpos = qpos + + # Create motion generator for smooth trajectory + motion_cfg = MotionGenCfg( + planner_cfg=ToppraPlannerCfg( + robot_uid=self.robot.uid, + ) + ) + self.motion_generator = MotionGenerator(cfg=motion_cfg) + + action_list = [] + + gripper_open = torch.tensor( + [0.05, 0.05], dtype=torch.float32, device=self.device + ) + gripper_close = torch.tensor( + [0.0, 0.0], dtype=torch.float32, device=self.device + ) + + # Convert waypoints to numpy for motion generator + qpos_waypoints_np = [q[0].cpu().numpy() for q in qpos_waypoints] + + # Generate trajectory segments + segments = [ + (0, 1, 30, gripper_open), + (1, 1, 10, gripper_close), + (1, 2, 30, gripper_close), + (2, 3, 40, gripper_close), + (3, 4, 30, gripper_close), + (4, 4, 10, gripper_open), + (4, 5, 30, gripper_open), + ] + + for start_idx, end_idx, num_steps, gripper_state in segments: + if start_idx == end_idx: + # Hold position (for gripper action) + qpos = qpos_waypoints[start_idx] + for _ in range(num_steps): + action = init_qpos.clone() + action[:, right_arm_ids] = qpos + action[:, right_eef_ids] = gripper_state.unsqueeze(0).expand( + self.num_envs, -1 + ) + action_list.append(action) + else: + # Generate smooth trajectory between waypoints + qpos_list = [qpos_waypoints_np[start_idx], qpos_waypoints_np[end_idx]] + + options = MotionGenOptions( + control_part="right_arm", + plan_opts=ToppraPlanOptions( + constraints={ + "velocity": 0.2, + "acceleration": 0.5, + }, + sample_method=TrajectorySampleMethod.QUANTITY, + sample_interval=num_steps, + ), + ) + # Joint space trajectory + target_states = [ + PlanState.single( + qpos=torch.as_tensor(qpos_waypoints_np[start_idx]), + move_type=MoveType.JOINT_MOVE, + ), + PlanState.single( + qpos=torch.as_tensor(qpos_waypoints_np[end_idx]), + move_type=MoveType.JOINT_MOVE, + ), + ] + plan_result = self.motion_generator.generate( + target_states=target_states, options=options + ) + + # Convert to torch and add to action list + for qpos_item in plan_result.positions[0]: + qpos = torch.as_tensor( + qpos_item, dtype=torch.float32, device=self.device + ) + + qpos = qpos.flatten() + + if qpos.shape[0] != len(right_arm_ids): + logger.log_warning( + f"Qpos shape mismatch: got {qpos.shape[0]}, expected {len(right_arm_ids)}" + ) + continue + + qpos = qpos.unsqueeze(0).expand(self.num_envs, -1) + + action = init_qpos.clone() + action[:, right_arm_ids] = qpos + action[:, right_eef_ids] = gripper_state.unsqueeze(0).expand( + self.num_envs, -1 + ) + action_list.append(action) + + logger.log_info( + f"Generated {len(action_list)} demo actions for stacking blocks" + ) + return action_list + + def is_task_success(self, **kwargs) -> torch.Tensor: + """Determine if the task is successfully completed. This is mainly used in the data generation process + of the imitation learning. + + The task is successful if: + 1. Block2 is stacked on top of Block1 + 2. Both blocks haven't fallen over + + Args: + **kwargs: Additional arguments for task-specific success criteria. + + Returns: + torch.Tensor: A boolean tensor indicating success for each environment in the batch. + """ + try: + block1 = self.sim.get_rigid_object("block_1") + block2 = self.sim.get_rigid_object("block_2") + except Exception as e: + logger.log_warning(f"Block1 or Block2 not found: {e}, returning False.") + return torch.zeros(self.num_envs, dtype=torch.bool, device=self.device) + + # Get poses + block1_pose = block1.get_local_pose(to_matrix=True) + block2_pose = block2.get_local_pose(to_matrix=True) + + # Extract positions + block1_pos = block1_pose[:, :3, 3] # (num_envs, 3) + block2_pos = block2_pose[:, :3, 3] + + # Check if blocks haven't fallen + block1_fallen = self._is_fall(block1_pose) + block2_fallen = self._is_fall(block2_pose) + + # Block2 should be on top of block1 + expected_block2_pos = torch.stack( + [ + block1_pos[:, 0], + block1_pos[:, 1], + block1_pos[:, 2] + 0.05, # block1 z + block height + ], + dim=1, + ) + + # Tolerance + eps = torch.tensor( + [0.025, 0.025, 0.012], dtype=torch.float32, device=self.device + ) + + # Check if block2 is within tolerance of expected position + position_diff = torch.abs(block2_pos - expected_block2_pos) # (num_envs, 3) + within_tolerance = torch.all( + position_diff < eps.unsqueeze(0), dim=1 + ) # (num_envs,) + + # Task succeeds if blocks are stacked correctly and haven't fallen + success = within_tolerance & ~block1_fallen & ~block2_fallen + + return success + + def _is_fall(self, pose: torch.Tensor) -> torch.Tensor: + # Extract z-axis from rotation matrix (last column, first 3 elements) + pose_rz = pose[:, :3, 2] + world_z_axis = torch.tensor([0, 0, 1], dtype=pose.dtype, device=pose.device) + + # Compute dot product for each batch element + dot_product = torch.sum(pose_rz * world_z_axis, dim=-1) # Shape: (batch_size,) + + # Clamp to avoid numerical issues with arccos + dot_product = torch.clamp(dot_product, -1.0, 1.0) + + # Compute angle and check if fallen + angle = torch.arccos(dot_product) + return angle >= torch.pi / 4 diff --git a/embodichain_tasks/embodichain_tasks/tableware/stack_cups.py b/embodichain_tasks/embodichain_tasks/tableware/stack_cups.py new file mode 100644 index 00000000..f8fbc44a --- /dev/null +++ b/embodichain_tasks/embodichain_tasks/tableware/stack_cups.py @@ -0,0 +1,118 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +import torch +import numpy as np + +from embodichain.lab.gym.envs import EmbodiedEnv, EmbodiedEnvCfg +from embodichain.lab.gym.utils.registration import register_env +from embodichain.utils import logger + +__all__ = ["StackCupsEnv"] + + +@register_env("StackCups-v1", max_episode_steps=600) +class StackCupsEnv(EmbodiedEnv): + def __init__(self, cfg: EmbodiedEnvCfg = None, **kwargs): + super().__init__(cfg, **kwargs) + + action_config = kwargs.get("action_config", None) + if action_config is not None: + self.action_config = action_config + + def is_task_success(self, **kwargs) -> torch.Tensor: + """Determine if the task is successfully completed. + + The task is successful if: + 1. Both cups haven't fallen over + 2. Cups are aligned in xy plane + 3. Top cup is only slightly higher than base cup + + Args: + **kwargs: Additional arguments for task-specific success criteria. + + Returns: + torch.Tensor: A boolean tensor indicating success for each environment in the batch. + """ + try: + cup1 = self.sim.get_rigid_object("cup_1") + cup2 = self.sim.get_rigid_object("cup_2") + except Exception as e: + logger.log_warning(f"Cups not found: {e}, returning False.") + return torch.zeros(self.num_envs, dtype=torch.bool, device=self.device) + + # Get cup poses + cup1_pose = cup1.get_local_pose(to_matrix=True) + cup2_pose = cup2.get_local_pose(to_matrix=True) + + # Extract positions + cup1_pos = cup1_pose[:, :3, 3] # (num_envs, 3) + cup2_pos = cup2_pose[:, :3, 3] # (num_envs, 3) + + # Check if cups haven't fallen + cup1_fallen = self._is_fall(cup1_pose) + cup2_fallen = self._is_fall(cup2_pose) + + # Determine which cup is lower (should be cup1) + cup1_is_lower = cup1_pos[:, 2] < cup2_pos[:, 2] + + # Use the lower cup as base, calculate expected position of top cup + base_cup_pos = torch.where( + cup1_is_lower.unsqueeze(1).expand_as(cup1_pos), cup1_pos, cup2_pos + ) + top_cup_pos = torch.where( + cup1_is_lower.unsqueeze(1).expand_as(cup2_pos), cup2_pos, cup1_pos + ) + + # Top cup should be slightly higher than base cup + cup_rim_offset = 0.015 + expected_top_pos = torch.stack( + [ + base_cup_pos[:, 0], + base_cup_pos[:, 1], + base_cup_pos[:, 2] + cup_rim_offset, + ], + dim=1, + ) + + # Tolerance + eps = torch.tensor([0.04, 0.04, 0.02], dtype=torch.float32, device=self.device) + + # Check if top cup is within tolerance of expected position + position_diff = torch.abs(top_cup_pos - expected_top_pos) # (num_envs, 3) + stacked_correctly = torch.all( + position_diff < eps.unsqueeze(0), dim=1 + ) # (num_envs,) + + # Task succeeds if cups are stacked correctly and haven't fallen + success = stacked_correctly & ~cup1_fallen & ~cup2_fallen + + return success + + def _is_fall(self, pose: torch.Tensor) -> torch.Tensor: + # Extract z-axis from rotation matrix (last column, first 3 elements) + pose_rz = pose[:, :3, 2] + world_z_axis = torch.tensor([0, 0, 1], dtype=pose.dtype, device=pose.device) + + # Compute dot product for each batch element + dot_product = torch.sum(pose_rz * world_z_axis, dim=-1) # Shape: (batch_size,) + + # Clamp to avoid numerical issues with arccos + dot_product = torch.clamp(dot_product, -1.0, 1.0) + + # Compute angle and check if fallen + angle = torch.arccos(dot_product) + return angle >= torch.pi / 4 diff --git a/embodichain_tasks/embodichain_tasks/utils/__init__.py b/embodichain_tasks/embodichain_tasks/utils/__init__.py new file mode 100644 index 00000000..355d915f --- /dev/null +++ b/embodichain_tasks/embodichain_tasks/utils/__init__.py @@ -0,0 +1,17 @@ +# ---------------------------------------------------------------------------- +# 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 diff --git a/embodichain_tasks/embodichain_tasks/utils/importer.py b/embodichain_tasks/embodichain_tasks/utils/importer.py new file mode 100644 index 00000000..bc3737f4 --- /dev/null +++ b/embodichain_tasks/embodichain_tasks/utils/importer.py @@ -0,0 +1,52 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Recursive sub-package importer for auto-registration of task environments. + +This follows the same pattern as IsaacLab's ``isaaclab_tasks`` — recursively +import every sub-package so that each task's ``__init__.py`` triggers its +``@register_env`` → ``gym.register()`` call chain. +""" + +from __future__ import annotations + +import importlib +import pkgutil + +__all__ = ["import_packages"] + + +def import_packages(package_name: str, blacklist: list[str] | None = None) -> None: + """Recursively import all sub-packages of *package_name*. + + Each imported sub-package executes its ``__init__.py``, which triggers + ``@register_env`` decorators that call ``gym.register()``. + + Args: + package_name: The fully-qualified package name (e.g. ``"embodichain_tasks"``). + blacklist: Sub-package names to skip (e.g. ``["utils"]``). + """ + blacklist = blacklist or [] + package = importlib.import_module(package_name) + + for _, name, is_pkg in pkgutil.walk_packages( + package.__path__, prefix=package_name + "." + ): + if any( + name.endswith("." + b) or name == package_name + "." + b for b in blacklist + ): + continue + importlib.import_module(name) diff --git a/embodichain_tasks/pyproject.toml b/embodichain_tasks/pyproject.toml new file mode 100644 index 00000000..7a7f2f58 --- /dev/null +++ b/embodichain_tasks/pyproject.toml @@ -0,0 +1,18 @@ +[build-system] +requires = ["setuptools>=61.0", "wheel"] +build-backend = "setuptools.build_meta" + +[project] +name = "embodichain_tasks" +version = "0.0.1" +description = "Official task environments for EmbodiChain" +requires-python = ">=3.10" +dependencies = [ + "embodichain", +] + +[project.entry-points."embodichain.tasks"] +"embodichain_tasks" = "embodichain_tasks" + +[tool.setuptools.packages.find] +include = ["embodichain_tasks*"]