So101 keyboard - #6876
Conversation
github-actions
Bot
added
documentation
Greptile SummaryThis PR adds a registered SO101 keyboard-typing environment with procedural keyboard assets, typing-state commands and curriculum logic, cross-backend scene presets, and a custom shared-encoder RSL-RL policy.
Confidence Score: 4/5The policy export failure should be fixed before merging because trained SO101 policies cannot be exported through the repository's standard RSL-RL workflow. The custom actor consumes named TensorDict groups during inference, while the standard exporter calls the copied actor with a single tensor and bypasses the custom split-input export adapters. Files Needing Attention: source/isaaclab_tasks/isaaclab_tasks/contrib/keyboard/agents/models.py Important Files Changed
Flowchart%%{init: {'theme': 'neutral'}}%%
flowchart LR
CLI[Train or play CLI] --> Gym[Isaac-Keyboard-SO101 registration]
Gym --> Env[ManagerBasedRLEnv]
Env --> Scene[SO101 robot and procedural keyboard]
Env --> MDP[Typing command, observations, rewards, resets]
MDP --> Wrapper[RSL-RL wrapper]
Wrapper --> PPO[SharedEncoderPPO]
PPO --> Actor[SharedEncoder actor]
PPO --> Critic[SharedEncoder critic]
Actor --> Export[Policy export workflow]
Reviews (1): Last reviewed commit: "clean up" | Re-trigger Greptile |
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
| latents = [self.encoders[group](obs[group]) for group in self.encoder_obs_groups] | ||
| if self.obs_groups: | ||
| latents.insert(0, super().get_latent(obs)) | ||
| return torch.cat(latents, dim=-1) |
There was a problem hiding this comment.
When exporting a trained SO101 policy through the standard RSL-RL JIT or ONNX workflow, the exporter passes the copied actor one concatenated tensor, while get_latent() indexes its input by named observation group. The exporter also bypasses as_jit() and as_onnx(), causing policy export to fail instead of producing a deployable model.
Knowledge Base Used: isaaclab_rl: RL Library Adapters
![isaaclab-review-bot[bot]](https://avatars.githubusercontent.com/u/157846462?s=60&v=4){kind=small}
There was a problem hiding this comment.
Isaac Lab Review Bot
The new SO-101 keyboard task has a coherent procedural asset and backend-specific articulation design, but the default-enabled replay curriculum mishandles the partitioned PhysX keyboard during snapshot/restore. The public task also depends on an internal development Nucleus asset, lacks the required isaaclab_tasks changelog fragment, and duplicates keyboard reset randomization.
- Design and architecture: The PhysX fixed-DOF versus Newton single-articulation split is appropriately isolated in the keyboard asset configuration and normalized through the command’s global slot mapping. However, the replay snapshot helpers bypass that mapping, index flattened articulation storage as though it were one row per environment, and depend directly on private InteractiveScene collections. The snapshot layer should use a partition-aware state mapping and public scene accessors.
- API: The Gym registration, lazy-exported modules, stubs, and environment documentation form a coherent new public surface. Before merge, the SO-101 robot asset must be hosted through a public asset root rather than the internal isaac-dev Nucleus server, and a minor changelog fragment must be added under source/isaaclab_tasks/changelog.d/.
- Implementation: For the default PhysX configuration, each environment has 18 keyboard articulation instances, but get_reset_state and set_reset_state use environment IDs directly against flattened articulation tensors and write APIs. This captures and restores the wrong instances, invalidating the enabled replay curriculum. The identical reset_root_state_uniform configuration in both EventCfg.reset_keyboard and typing.reset.pre_solve_reset is also redundant because the later pre-solve write or snapshot restore replaces the event result.
Significant concerns. Posted 5 actionable findings inline.
Automated review; human maintainers own approval decisions.
| states: list[Tensor] = [] | ||
| for name, articulation in env.scene._articulations.items(): | ||
| if name in reset_assets: | ||
| root_state = wp.to_torch(articulation.data.root_state_w)[env_ids] |
There was a problem hiding this comment.
🔴 Critical · Implementation — Snapshot indexes flattened articulation data by env id
For the default PhysX preset the keyboard exposes 18 articulation instances per environment (parts/part_*), which is why LetterTypingCommand builds explicit _inst_idx maps. Here root_state_w, joint_pos, and joint_vel are indexed directly with env_ids, so snapshots capture one arbitrary partition and set_reset_state writes those ids back as instance indices, corrupting other environments' keys. Gather and restore all (env, partition) instances.
| # Register Gym environments. | ||
| ## | ||
|
|
||
| gym.register( |
There was a problem hiding this comment.
🟡 Warning · Api — Missing isaaclab_tasks changelog fragment
This PR adds a registered environment and new public modules under source/isaaclab_tasks and also edits core/lift, but no fragment was added under source/isaaclab_tasks/changelog.d/. Repository rules require one fragment per touched package (CI compiles CHANGELOG.rst and the version bump from these). Add e.g. so101-keyboard.minor.rst with an Added entry for the new task.
| from . import mdp | ||
| from .keyboards import TYPING_KEYBOARD_POOL | ||
|
|
||
| _SO101_USD_PATH = "omniverse://isaac-dev.ov.nvidia.com/Isaac/IsaacLab/Robots/SO101/so101.usda" |
There was a problem hiding this comment.
🟡 Warning · Implementation — Robot USD points at internal dev Nucleus
_SO101_USD_PATH hardcodes omniverse://isaac-dev.ov.nvidia.com/..., while the sky light in the same file resolves through ISAAC_NUCLEUS_DIR. Since Isaac-Keyboard-SO101 is registered and listed in environments.rst, external users cannot spawn the robot at all. Publish the asset under the public nucleus root and build the path from ISAAC_NUCLEUS_DIR.
| def get_reset_state(env, env_ids: Tensor, reset_assets: Sequence[str], is_relative: bool = False) -> Tensor: | ||
| """Read and concatenate reset-state slices for scene assets.""" | ||
| states: list[Tensor] = [] | ||
| for name, articulation in env.scene._articulations.items(): |
There was a problem hiding this comment.
🔵 Suggestion · Design Architecture — Reset helpers use private scene collections
Both helpers iterate env.scene._articulations and env.scene._rigid_objects, private members of InteractiveScene from another package, while equivalent public mappings are available. Any rename of the internal storage silently breaks the curriculum snapshot/restore path. Use the public scene accessors instead.
| class EventCfg: | ||
| """Reset-mode events (shared by all physics backends).""" | ||
|
|
||
| reset_keyboard = EventTerm( |
There was a problem hiding this comment.
🔵 Suggestion · Implementation — Keyboard randomization configured twice
EventCfg.reset_keyboard and typing.reset.pre_solve_reset apply reset_root_state_uniform to the keyboard with identical ranges. Reset-mode events run before command_manager.reset(), so the event's write is always overwritten (by the pre-solve reset on the normal path, by the snapshot restore on the buffer path). The pre_solve_reset docstring itself says to keep only one owner; drop the duplicate.
ae238a4 to
174b84d
Compare
AntoineRichard
left a comment
AntoineRichard
left a comment
There was a problem hiding this comment.
AI-generated review — request changes
I reviewed the current head (ba2bb025e9) with emphasis on functional comments, duplication, test value, documentation, and strict showroom-quality hygiene. The new inline comments below are additional findings; I did not duplicate the existing precise threads.
The following existing unresolved findings remain blocking on this head: the standard RSL-RL export path bypasses the custom adapters, partitioned PhysX reset snapshots index flattened articulation state incorrectly, the robot uses an internal Nucleus URL, reset-state helpers depend on private scene collections, and keyboard randomization is configured twice. The earlier changelog finding is addressed by so101-keyboard.minor.rst and can be resolved.
Tests
This 5.3k-line change currently adds or modifies no test files, despite the PR description claiming focused configuration and architecture tests. The existing contributed-environment smoke test should not be duplicated; it is already useful and CI shows it failing for this task with FileNotFoundError on the internal SO-101 USD. Add a small consolidated suite with distinct contracts:
- one CPU model/export test that exercises the exact standard JIT and ONNX entry points, not only
as_jit()/as_onnx()in isolation; - one parametrized typing-state test covering simultaneous presses, backspace, buffer-full behavior, padding, and partial
env_ids; - one reset-state round-trip test covering both single-articulation Newton layout and partitioned PhysX layout;
- one parametrized pure geometry/pool test only if the unused general generator remains.
Use shared fixtures and parametrization. Do not add tests that merely restate constants, inspect config fields, or repeat task registration/instantiation already covered by the generic smoke test. Each test should fail for a named behavioral regression.
Verification
CUDA_VISIBLE_DEVICES='' ./isaaclab.sh -f passes all hooks. I did not run local simulator/GPU tests because the machine is running benchmarks and GPU use was explicitly excluded. Current CI has failures, including the task smoke failure above. GitHub also reports the PR as non-mergeable (mergeable_state: dirty).
| @@ -127,6 +128,7 @@ | |||
| [/Reorient-Franka/, "tasks/manipulation/franka_lift.jpg"], | |||
| [/Reorient-KukaAllegro/, "tasks/manipulation/kuka_allegro_reorient.jpg"], | |||
| [/Shadow-Handover/, "tasks/manipulation/shadow_hand_over.jpg"], | |||
| [/Keyboard-SO101/, "tasks/manipulation/so101_keyboard.jpg"], | |||
There was a problem hiding this comment.
AI-generated review — Important (documentation): This rule points to docs/source/_static/tasks/manipulation/so101_keyboard.jpg, but that file is absent from both the PR and the base tree. The browser will request a missing _images/so101_keyboard.jpg; the default image only applies when no regex matches and does not handle this 404. Add the licensed screenshot in this PR, or remove this rule until the asset exists.
| @@ -86,6 +86,7 @@ class ObjectCfg(PresetCfg): | |||
| mass_props=sim_utils.MassPropertiesCfg(mass=0.2), | |||
| ) | |||
| default = shapes | |||
| ovphysx = cube | |||
There was a problem hiding this comment.
AI-generated review — Important (API/scope): This unrelated preset move is coupled to deleting the top-level KukaAllegroObjectCfg class in kuka_allegro_env_cfg.py. That is a public symbol removal without a prior deprecation, contrary to the repository's breaking-change rule. Revert the three core/lift edits from this keyboard PR. If the consolidation is needed, handle it separately with deprecation and focused regression coverage.
| slot_labels: tuple[str, ...] | ||
|
|
||
|
|
||
| _BASE_SPAWNER = KeyboardSpawnerCfg(family="ansi_full", topology_mode="exact", partition_dof=_PARTITION_DOF) |
There was a problem hiding this comment.
AI-generated review — Important (scope/documentation): Every task variant derives from family="ansi_full", so the environment does not vary layouts as the PR description claims. The other layout families, topology modes, and tail policy add a large unexercised surface. Either make layout variation real while preserving the task's stable slot/DOF contract and test it, or correct the description and remove/split the unused generalization.
| max_len=5, | ||
| # letter_full shows both target and typed letters; letter_left shows only the remaining target | ||
| command_mode="letter_full", | ||
| typeable_slots=tuple( |
There was a problem hiding this comment.
AI-generated review — Important (behavioral contract): This excludes only Backspace, so targets include Ctrl, Shift, F-keys, arrows, media keys, and punctuation. That is an arbitrary key sequence, not the letter / target word behavior documented throughout the task. Either restrict targets to the intended alphabetic slots (and test the exact set), or consistently rename the command and documentation to key-sequence typing.
| UsdGeom.Scope.Define(stage, part_keys_scope) | ||
| UsdGeom.Scope.Define(stage, part_joints_scope) | ||
|
|
||
| part_base_prim = stage.GetPrimAtPath(part_base_path) |
There was a problem hiding this comment.
AI-generated review — Moderate (duplication): The rigid-body, mass, tiny-inertia, and world-fixed-joint setup duplicates the single-articulation block at lines 127–137. Extract one helper parameterized by base path, joint scope, and mass; keep only the partition-specific policy here. This is physics-authoring code, so a single owner for those authored properties is preferable to parallel branches.
| self._buffer_built = True | ||
| self._log_buffer_stats() | ||
|
|
||
| def _log_buffer_stats(self): |
There was a problem hiding this comment.
AI-generated review — Moderate (production hygiene): This unconditionally runs roughly 75 lines of diagnostic reductions and print() calls during the first reset, including multiple device-to-host synchronizations and verbose histograms. Keep essential metrics through the repository logging/metrics facilities and gate any detailed coverage dump behind an explicit debug option. Production task startup should not emit a debug report by default.
| return (command.distance == 0).float() | ||
|
|
||
|
|
||
| class reach_key(ManagerTermBase): |
There was a problem hiding this comment.
AI-generated review — Moderate (dead public surface): Neither reach_key nor typing_mistake is referenced by this environment configuration, yet both are exported from mdp/__init__.pyi and neither is tested. Remove these unused terms and exports, or wire them into a documented behavior with distinct tests. Shipping unused MDP terms expands the public surface and, for typing_mistake, contradicts the configured backspace-recovery curriculum.
| ) | ||
|
|
||
|
|
||
| def is_identity_quat(quat_xyzw: tuple[float, float, float, float]) -> bool: |
There was a problem hiding this comment.
AI-generated review — Moderate (dead code): is_identity_quat has no callers; quat_to_matrix and mat3_rotate below are also unused. Remove all three rather than carrying an untested parallel math utility layer. Keep only the quaternion helpers actually needed by keyboard generation.
| a policy trained with ``letter_length=(1, 2)`` on single letters by keeping ``max_len=2`` and setting | ||
| ``letter_length=(1, 1)``.""" | ||
|
|
||
| command_mode: str = "letter_full" |
There was a problem hiding this comment.
AI-generated review — Moderate (type design): Only two values are valid, but the type is unrestricted str, and any typo silently takes the letter_full branch. Use Literal["letter_full", "letter_left"] and validate at construction/config boundaries. Include one negative test that proves invalid modes fail clearly.
AntoineRichard
left a comment
AntoineRichard
left a comment
There was a problem hiding this comment.
AI-generated review — performance follow-up
I audited the new MDP asset-write paths for reusable indices. The reset-state helpers already receive tensor env_ids; the reset IK joint selection is the concrete hot-path violation and is called repeatedly inside the solve loop.
| if cfg.reset.ik is not None: | ||
| ik_cfg = cfg.reset.ik | ||
| self.robot: Articulation = env.scene[cfg.asset_name] | ||
| self._ik_joint_ids, _ = self.robot.find_joints(list(ik_cfg.joint_names)) |
There was a problem hiding this comment.
AI-generated review — important (performance): cache device-native joint indices
find_joints() returns a Python list here, and _solve_reset_pose() passes it to write_joint_position_to_sim_index() on every IK iteration. The articulation backends resolve a list by constructing a device index array each call, adding allocation/conversion overhead inside this reset hot loop. Request the cached proxy (as_proxy=True), store its device-native tensor once (for example, .torch because these IDs also index torch data), and derive _ik_jacobi_joint_ids as a tensor as well. Asset writes in MDPs should receive pre-allocated tensor/Warp indices—or None/a slice for the full selection—rather than Python lists.
AntoineRichard
left a comment
AntoineRichard
left a comment
There was a problem hiding this comment.
AI-generated review — simplification follow-up
There is substantial room to reduce this PR without weakening the task. The largest cut remains the earlier inline finding to remove the unused layout/topology framework: the runtime pool is always a 108-key ansi_full layout. The earlier comments also cover deleting unused MDP terms/math helpers, consolidating duplicated USD authoring and prefix logic, removing the verbose startup report, shortening configuration prose, and eliminating the duplicated keyboard reset.
The additional inline findings below identify three non-overlapping reductions: reuse the existing reset-state implementation, remove an unused manifest surface, and avoid eagerly constructing every style archetype. Together with the earlier findings, this should remove several hundred lines and reduce the amount of behavior that needs tests.
ba2bb02 to
c342f75
Compare
hujc7
left a comment
hujc7
left a comment
There was a problem hiding this comment.
Agent review: three findings on head f100f4f, checked against the existing threads so they do not restate them.
The main one is that deleting mdp/states.py in favour of the shared core/lift helper looks like it closes the partitioned-PhysX reset thread, but the incorrect indexing came along with the move — details inline. Two sub-findings on that thread (private scene collections, deprecated root-state APIs) genuinely are fixed and can be resolved.
The other two are small: the default physics preset contradicts this file's own guidance, and a reward config class carries the reorient task's name.
|
|
||
| from isaaclab_tasks.core.lift.mdp.events import SuccessMonitor | ||
| from isaaclab_tasks.core.lift.mdp.events_cfg import SuccessMonitorCfg | ||
| from isaaclab_tasks.core.lift.mdp.utils import get_reset_state, set_reset_state |
There was a problem hiding this comment.
Agent found bug: the partitioned-reset fix moved the bug rather than fixing it. Deleting mdp/states.py in favour of the shared helper resolved two of the sub-findings, but not the correctness one.
core/lift/mdp/utils.py::get_reset_state reads articulation.data.joint_pos.torch[env_ids], and joint_pos is (num_instances, num_joints) — not (num_envs, ...) (isaaclab/assets/articulation/base_articulation_data.py:925). On the default isaacsim_physx preset the keyboard has 18 articulation roots per env, as this PR documents at so101_env_cfg.py:40-42 and as partition_dof=6 over 108 slots produces.
So env_ids = [0..k) selects part_0..part_{k-1} of environment 0, and each snapshot stores 13 + 2*6 = 25 floats — 6 of 108 key joints. set_reset_state writes them back with the same conflation. The shapes stay self-consistent, so nothing raises.
reset.enabled=True (so101_env_cfg.py:154) with reset.ik set, so the replay curriculum is on by default on the preset this breaks. Newton (1 instance/env, 108 DOF) is unaffected, which is why a Newton-only training run does not surface it.
Two sub-findings are fixed by the move, and those threads can be resolved: the shared helper uses the public scene.articulations accessor, and it uses root_link_pose_w / write_root_link_pose_to_sim_index instead of the deprecated root_state_w / write_root_state_to_sim.
What remains is the partition-aware mapping requested on the old states.py:13. LetterTypingCommand.__init__ already builds _inst_idx / _col_idx for exactly this layout; the shared helper needs the same awareness, or an explicit per-asset instance-index argument, before either caller can round-trip a multi-root articulation.
- Reset-state round-trip test over both layouts — Newton single-articulation and partitioned PhysX — asserting
restore(capture(s)) == sfor a non-trivialenv_idssubset.
| debug_mode=False, | ||
| ) | ||
| physx = PhysxAutoCfg(isaacsim_physx=isaacsim_physx) | ||
| default = isaacsim_physx |
There was a problem hiding this comment.
Agent suggestion: the default physics preset is the one this file advises against, and the one the partitioned-reset issue breaks.
:288 states "physx is usable but extremely slow for this task and is only for evaluation purposes / for training please only use newton_mjwarp", yet :315 sets default = isaacsim_physx. Anyone omitting physics= gets the slow path, which is also the 18-roots-per-env layout the reset snapshot mishandles.
Worth noting the fix is not a one-line flip. resolve_presets picks per PresetCfg by selected name and otherwise falls back to that preset's own default (isaaclab_tasks/utils/hydra.py:419-420). With no physics= on the command line, KeyboardAssetCfg would still resolve to its default at :47 — the partitioned PhysX articulation — while the sim ran Newton. Per this class's own docstring, Newton reads only the first articulation per env, so the other 17 key banks would be invisible. Changing the default means moving KeyboardAssetCfg.default to the single-articulation variant at the same time.
If PhysX has to remain the default for checkpoint parity, that is a reasonable call — but please state it in the PhysicsCfg docstring. As written, the comment reads as advice against the default the class sets.
|
|
||
|
|
||
| @configclass | ||
| class SO101ReorientRewardCfg: |
There was a problem hiding this comment.
Agent review: SO101ReorientRewardCfg names the reorient task, not this one. It reads as a copy-paste from the lift/reorient family and will misdirect anyone grepping for this task's reward config.
| class SO101ReorientRewardCfg: | |
| class SO101TypingRewardCfg: |
The reference in SO101KeyboardEnvCfg.rewards needs the matching rename.
StafaH
left a comment
StafaH
left a comment
There was a problem hiding this comment.
LGTM. The entire keyboard setup section looks extremely complex though, but I don't have any ideas on how to clean it up.
|
|
||
| __all__ = ["KeyboardSpawnerCfg", "TYPING_KEYBOARD_POOL"] | ||
|
|
||
| from .keyboard_gen_cfg import KeyboardSpawnerCfg |
There was a problem hiding this comment.
absolute imports here and elsewhere
7808b77 to
95d35f0
Compare
95d35f0 to
4ac1354
Compare
5 participants
Description
Adds the
IsaacContrib-Keyboard-SO101manager-based reinforcement-learning task underisaaclab_tasks.contrib.keyboard. The task trains an SO-101 arm to type target words while preserving the training and evaluation behavior validated by the accompanying checkpoint.The change includes:
isaacsim_physx, the default), automatic PhysX selection (physx), and Newton MJWarp (newton_mjwarp).Implementation reference: Emergent Dexterity via Diverse Resets and Large-Scale Reinforcement Learning (OmniReset).
Asset dependency
The SO-101 USD is hosted separately and is intentionally not included in this PR. Before merge, it should be published under the official IsaacLab Nucleus
Robots/SO101/directory and resolved throughISAAC_NUCLEUS_DIR.Training and evaluation
The policy can be evaluated with:
Manual checkpoint playback validates that the policy solves the typing task reliably.
Type of change
Screenshots
A task screenshot still needs to be attached and added to the environment catalog at:
docs/source/_static/tasks/manipulation/so101_keyboard.jpgChecklist
pre-commitchecks with./isaaclab.sh --formatsource/<pkg>/changelog.d/for every touched package (do not editCHANGELOG.rstor bumpextension.toml— CI handles that)CONTRIBUTORS.mdor my name already exists there