So101 keyboard (#6876)

# Description

Adds the `IsaacContrib-Keyboard-SO101` manager-based
reinforcement-learning task under `isaaclab_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:

- Procedurally generated, logo-agnostic keyboard families with varied
layouts, geometry, materials, and articulated keys. Both PhysX and
Newton use the full keyboard pool rather than a single fixed keyboard.
- Typing commands, observations, rewards, terminations, reset-state
replay, success-rate curriculum sampling, and visualization markers.
- A single RSL-RL actor/critic configuration with a shared MLP encoder
for the one-hot typing command.
- Physics presets for concrete Isaac Sim PhysX (`isaacsim_physx`, the
default), automatic PhysX selection (`physx`), and Newton MJWarp
(`newton_mjwarp`).
- Focused configuration and architecture tests, plus environment-catalog
documentation.

Implementation reference: [Emergent Dexterity via Diverse Resets and
Large-Scale Reinforcement Learning
(OmniReset)](https://arxiv.org/abs/2603.15789).

### 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 through
`ISAAC_NUCLEUS_DIR`.

### Training and evaluation

The policy can be evaluated with:

```bash
./isaaclab.sh play --task=IsaacContrib-Keyboard-SO101 --num_envs=32 physics=newton_mjwarp --rl_library=rsl_rl
```

Manual checkpoint playback validates that the policy solves the typing
task reliably.

## Type of change

- New feature (non-breaking change which adds functionality)

## Screenshots

A task screenshot still needs to be attached and added to the
environment catalog at:

`docs/source/_static/tasks/manipulation/so101_keyboard.jpg`

## Checklist

- [ ] I have read and understood the [contribution
guidelines](https://isaac-sim.github.io/IsaacLab/main/source/refs/contributing.html)
- [x] I have run the [`pre-commit` checks](https://pre-commit.com/) with
`./isaaclab.sh --format`
- [x] I have made corresponding changes to the documentation
- [ ] My changes generate no new warnings
- [x] I have added tests that prove my feature works
- [ ] I have added a changelog fragment under
`source/<pkg>/changelog.d/` for every touched package (do **not** edit
`CHANGELOG.rst` or bump `extension.toml` — CI handles that)
- [x] I have added my name to the `CONTRIBUTORS.md` or my name already
exists there

Author: ooctipus

docs/source/_static/css/environment-browser.js

@@ -21,6 +21,7 @@
             ["Isaac-Fourbar-Pole-Swingup", "rsl_rl", "newton_kamino", "", ""],
             ["Isaac-Humanoid-Direct", "rl_games,rsl_rl,skrl", "isaacsim_physx,newton_mjwarp,ovphysx", "", ""],
             ["Isaac-Humanoid", "rl_games,rsl_rl,skrl,sb3", "isaacsim_physx,newton_mjwarp,ovphysx", "", ""],
+            ["IsaacContrib-Keyboard-SO101", "rsl_rl", "isaacsim_physx,newton_mjwarp,physx", "", ""],
             ["Isaac-Lift-Cable-Franka", "rsl_rl", "newton_mjwarp_vbd_proxy", "", "ik,joint"],
             ["Isaac-Lift-Cable-Franka-Camera", "rsl_rl", "newton_mjwarp_vbd_proxy", "isaacsim_rtx,newton_renderer,ovrtx", "ik,joint"],
             ["Isaac-Lift-Cloth-Franka", "rsl_rl", "isaacsim_physx,newton_mjwarp_vbd_proxy", "", "ik,joint"],
@@ -93,7 +94,7 @@
         if (task.includes("Velocity")) {
             return "locomotion";
         }
-        if (/Lift|Reach|Reorient|Drawer|Handover/.test(task)) {
+        if (/Lift|Reach|Reorient|Drawer|Handover|Keyboard/.test(task)) {
             return "manipulation";
         }
         return "classic";
@@ -130,6 +131,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"],
             [/AnymalB/, "tasks/locomotion/anymal_b_flat.jpg"],
             [/AnymalC/, "tasks/locomotion/anymal_c_flat.jpg"],
             [/AnymalD/, "tasks/locomotion/anymal_d_flat.jpg"],

docs/source/overview/environments.rst

@@ -1413,6 +1413,10 @@ including disabling runtime perturbations used for training.
       - Direct
       - **skrl** (AMP)
       -
+    * - IsaacContrib-Keyboard-SO101
+      - Manager Based
+      - **rsl_rl** (PPO)
+      - **physics=** ``isaacsim_physx``, ``newton_mjwarp``
     * - IsaacContrib-Lift-Cube-Franka
       - Manager Based
       - **rl_games** (PPO), **rsl_rl** (PPO), **skrl** (PPO), **sb3** (PPO)

source/isaaclab_tasks/changelog.d/so101-keyboard.minor.rst

@@ -0,0 +1,16 @@
+Added
+^^^^^
+
+* Added the ``IsaacContrib-Keyboard-SO101`` task with procedural keyboards and PhysX and Newton support.
+* Added :func:`~isaaclab_tasks.core.lift.mdp.utils.env_instance_rows`, mapping each environment to the
+  instance rows of an asset that has several roots per environment.
+
+Changed
+^^^^^^^
+
+* Changed :func:`~isaaclab_tasks.core.lift.mdp.utils.get_reset_state` and
+  :func:`~isaaclab_tasks.core.lift.mdp.utils.set_reset_state` to cover every instance an asset has in an
+  environment instead of reading and writing the instance row that matches the environment id, which
+  addressed the wrong instances for assets partitioned into several roots per environment. Serialized
+  states are unchanged for assets with one instance per environment; for partitioned assets each slice
+  now repeats per instance, so states captured before this change cannot be restored with it.

source/isaaclab_tasks/isaaclab_tasks/contrib/keyboard/__init__.py

@@ -0,0 +1,35 @@
+# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md).
+# All rights reserved.
+#
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""SO-101 keyboard-typing environments.
+
+The MDP implementation is based primarily on OmniReset:
+
+@inproceedings{yin2026emergent,
+  title={Emergent Dexterity via Diverse Resets and Large-Scale Reinforcement Learning},
+  author={Patrick Yin and Tyler Westenbroek and Zhengyu Zhang and Joshua Tran and Ignacio Dagnino and Eeshani Shilamkar and Numfor Mbiziwo-Tiapo and Simran Bagaria and Xinlei Liu and Galen Mullins and Andrey Kolobov and Abhishek Gupta},
+  booktitle={The Fourteenth International Conference on Learning Representations},
+  year={2026},
+  url={https://arxiv.org/abs/2603.15789},
+}
+"""
+
+import gymnasium as gym
+
+from . import agents
+
+##
+# Register Gym environments.
+##
+
+gym.register(
+    id="IsaacContrib-Keyboard-SO101",
+    entry_point="isaaclab.envs:ManagerBasedRLEnv",
+    disable_env_checker=True,
+    kwargs={
+        "env_cfg_entry_point": f"{__name__}.so101_env_cfg:SO101KeyboardEnvCfg",
+        "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:SO101PPORunnerCfg",
+    },
+)

source/isaaclab_tasks/isaaclab_tasks/contrib/keyboard/agents/__init__.py

@@ -0,0 +1,4 @@
+# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md).
+# All rights reserved.
+#
+# SPDX-License-Identifier: BSD-3-Clause

source/isaaclab_tasks/isaaclab_tasks/contrib/keyboard/agents/models.py

@@ -0,0 +1,205 @@
+# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md).
+# All rights reserved.
+#
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""RSL-RL models for the SO-101 keyboard-typing task."""
+
+from __future__ import annotations
+
+import copy
+from typing import Any
+
+import torch
+import torch.nn as nn
+from rsl_rl.algorithms import PPO
+from rsl_rl.models import MLPModel
+from rsl_rl.modules import MLP, HiddenState
+from rsl_rl.storage import RolloutStorage
+from tensordict import TensorDict
+
+
+class SharedEncoderMLPModel(MLPModel):
+    """Encode selected 1-D observation groups before the MLP head."""
+
+    def __init__(
+        self,
+        obs: TensorDict,
+        obs_groups: dict[str, list[str]],
+        obs_set: str,
+        output_dim: int,
+        hidden_dims: tuple[int, ...] | list[int] = (256, 256, 256),
+        activation: str = "elu",
+        obs_normalization: bool = False,
+        distribution_cfg: dict | None = None,
+        encoder_cfg: dict[str, dict[str, Any]] | None = None,
+    ) -> None:
+        """Initialize the encoded-observation MLP model."""
+        if not encoder_cfg:
+            raise ValueError("At least one encoder configuration must be provided.")
+
+        active_obs_groups = obs_groups[obs_set]
+        encoder_keys = set(encoder_cfg)
+        if not encoder_keys.issubset(active_obs_groups):
+            invalid_groups = sorted(encoder_keys - set(active_obs_groups))
+            raise ValueError(
+                f"The encoder observation groups {invalid_groups} are not part of the '{obs_set}' observation groups"
+                f" {active_obs_groups}."
+            )
+
+        self.encoder_obs_groups = [group for group in active_obs_groups if group in encoder_keys]
+        self.encoder_input_dims = []
+        for obs_group in self.encoder_obs_groups:
+            if len(obs[obs_group].shape) != 2:
+                raise ValueError(
+                    f"The MLP encoders only support 1D observations, got shape {obs[obs_group].shape} for"
+                    f" '{obs_group}'."
+                )
+            self.encoder_input_dims.append(obs[obs_group].shape[-1])
+
+        encoders = {}
+        self.encoder_latent_dim = 0
+        for obs_group, input_dim in zip(self.encoder_obs_groups, self.encoder_input_dims):
+            group_cfg = dict(encoder_cfg[obs_group])
+            latent_dim = group_cfg["latent_dim"]
+            encoders[obs_group] = MLP(
+                input_dim=input_dim,
+                output_dim=latent_dim,
+                hidden_dims=group_cfg["hidden_dims"],
+                activation=group_cfg.get("activation", "elu"),
+                last_activation=group_cfg.get("last_activation"),
+            )
+            self.encoder_latent_dim += latent_dim
+
+        super().__init__(
+            obs,
+            obs_groups,
+            obs_set,
+            output_dim,
+            hidden_dims,
+            activation,
+            obs_normalization,
+            distribution_cfg,
+        )
+        self.encoders = nn.ModuleDict(encoders)
+
+    def get_latent(
+        self, obs: TensorDict, masks: torch.Tensor | None = None, hidden_state: HiddenState = None
+    ) -> torch.Tensor:
+        """Build the model latent from raw observations and encoded groups."""
+        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)
+
+    def update_normalization(self, obs: TensorDict) -> None:
+        """Update normalization statistics of non-encoded observation groups."""
+        if self.obs_groups:
+            super().update_normalization(obs)
+
+    def as_jit(self) -> nn.Module:
+        """Return a version of the model compatible with Torch JIT export."""
+        return _TorchSharedEncoderModel(self)
+
+    def as_onnx(self, verbose: bool = False) -> nn.Module:
+        """Return a version of the model compatible with ONNX export."""
+        return _OnnxSharedEncoderModel(self, verbose)
+
+    def _get_obs_dim(self, obs: TensorDict, obs_groups: dict[str, list[str]], obs_set: str) -> tuple[list[str], int]:
+        """Select non-encoded observation groups and compute their total dimension."""
+        active_obs_groups = obs_groups[obs_set]
+        raw_obs_groups = []
+        obs_dim = 0
+        for obs_group in active_obs_groups:
+            if len(obs[obs_group].shape) != 2:
+                raise ValueError(
+                    f"The MLP model only supports 1D observations, got shape {obs[obs_group].shape} for '{obs_group}'."
+                )
+            if obs_group not in self.encoder_obs_groups:
+                raw_obs_groups.append(obs_group)
+                obs_dim += obs[obs_group].shape[-1]
+        return raw_obs_groups, obs_dim
+
+    def _get_latent_dim(self) -> int:
+        """Return the latent dimensionality consumed by the MLP head."""
+        return self.obs_dim + self.encoder_latent_dim
+
+
+class SharedEncoderPPO(PPO):
+    """Share the actor's observation encoders with the critic."""
+
+    def __init__(self, actor: MLPModel, critic: MLPModel, storage: RolloutStorage, **kwargs: Any) -> None:
+        """Replace the critic's encoders before PPO registers optimizer parameters."""
+        if not isinstance(actor, SharedEncoderMLPModel) or not isinstance(critic, SharedEncoderMLPModel):
+            raise TypeError("SharedEncoderPPO requires SharedEncoderMLPModel actor and critic models.")
+        if (
+            actor.encoder_obs_groups != critic.encoder_obs_groups
+            or actor.encoder_latent_dim != critic.encoder_latent_dim
+        ):
+            raise ValueError("The actor and critic encoder configurations must match.")
+
+        # The actor owns the shared modules so optimizer, checkpoint, and gradient traversal see them once.
+        del critic.encoders
+        object.__setattr__(critic, "encoders", actor.encoders)
+        super().__init__(actor, critic, storage, **kwargs)
+
+
+class _TorchSharedEncoderModel(nn.Module):
+    """Exportable shared-encoder model for TorchScript."""
+
+    def __init__(self, model: SharedEncoderMLPModel) -> None:
+        """Create a TorchScript-compatible model copy."""
+        super().__init__()
+        self.obs_normalizer = copy.deepcopy(model.obs_normalizer)
+        self.encoders = nn.ModuleList([copy.deepcopy(model.encoders[g]) for g in model.encoder_obs_groups])
+        self.mlp = copy.deepcopy(model.mlp)
+        if model.distribution is not None:
+            self.deterministic_output = model.distribution.as_deterministic_output_module()
+        else:
+            self.deterministic_output = nn.Identity()
+
+    def forward(self, obs_raw: torch.Tensor, obs_encoded: list[torch.Tensor]) -> torch.Tensor:
+        """Run deterministic inference from raw and encoded-group inputs."""
+        latents = [self.obs_normalizer(obs_raw)]
+        for i, encoder in enumerate(self.encoders):
+            latents.append(encoder(obs_encoded[i]))
+        return self.deterministic_output(self.mlp(torch.cat(latents, dim=-1)))
+
+    @torch.jit.export
+    def reset(self) -> None:
+        """Reset recurrent export state."""
+        pass
+
+
+class _OnnxSharedEncoderModel(_TorchSharedEncoderModel):
+    """Exportable shared-encoder model for ONNX."""
+
+    is_recurrent: bool = False
+
+    def __init__(self, model: SharedEncoderMLPModel, verbose: bool) -> None:
+        """Create an ONNX-compatible model copy."""
+        super().__init__(model)
+        self.verbose = verbose
+        self.encoder_obs_groups = list(model.encoder_obs_groups)
+        self.encoder_input_dims = list(model.encoder_input_dims)
+        self.obs_dim_raw = model.obs_dim
+
+    def forward(self, obs: torch.Tensor, *obs_encoded: torch.Tensor) -> torch.Tensor:
+        """Run deterministic inference for ONNX export."""
+        return super().forward(obs, list(obs_encoded))
+
+    def get_dummy_inputs(self) -> tuple[torch.Tensor, ...]:
+        """Return representative dummy inputs for ONNX tracing."""
+        dummy_raw = torch.zeros(1, self.obs_dim_raw)
+        dummy_encoded = [torch.zeros(1, dim) for dim in self.encoder_input_dims]
+        return (dummy_raw, *dummy_encoded)
+
+    @property
+    def input_names(self) -> list[str]:
+        """Return ONNX input tensor names."""
+        return ["obs", *self.encoder_obs_groups]
+
+    @property
+    def output_names(self) -> list[str]:
+        """Return ONNX output tensor names."""
+        return ["actions"]

source/isaaclab_tasks/isaaclab_tasks/contrib/keyboard/agents/rsl_rl_ppo_cfg.py

@@ -0,0 +1,95 @@
+# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md).
+# All rights reserved.
+#
+# SPDX-License-Identifier: BSD-3-Clause
+
+from dataclasses import MISSING
+
+from isaaclab.utils.configclass import configclass
+
+from isaaclab_rl.rsl_rl import (
+    RslRlMLPModelCfg,
+    RslRlOnPolicyRunnerCfg,
+    RslRlPpoAlgorithmCfg,
+)
+
+##
+# Custom model configurations.
+##
+
+
+@configclass
+class SharedEncoderMLPModelCfg(RslRlMLPModelCfg):
+    """Configuration for the shared-encoder MLP model."""
+
+    class_name: str = "isaaclab_tasks.contrib.keyboard.agents.models:SharedEncoderMLPModel"
+    """The model class name. Defaults to :class:`~.models.SharedEncoderMLPModel`."""
+
+    @configclass
+    class EncoderCfg:
+        """Configuration for the MLP encoder of a single observation group."""
+
+        hidden_dims: list[int] = MISSING
+        """The hidden dimensions of the encoder MLP."""
+
+        latent_dim: int = MISSING
+        """The dimension of the encoder output latent."""
+
+        activation: str = "elu"
+        """The activation function of the encoder MLP. Defaults to elu."""
+
+        last_activation: str | None = "elu"
+        """The activation applied to the encoder output latent. Defaults to elu.
+
+        If None, the latent is the output of the last linear layer.
+        """
+
+    encoder_cfg: dict[str, EncoderCfg] = MISSING
+    """Mapping from observation group name to the MLP encoder configuration for that group."""
+
+
+##
+# Runner configuration.
+##
+
+
+@configclass
+class SO101PPORunnerCfg(RslRlOnPolicyRunnerCfg):
+    """RSL-RL configuration for SO-101 keyboard typing with shared observation encoders."""
+
+    num_steps_per_env = 32
+    max_iterations = 15000
+    save_interval = 250
+    experiment_name = "so101_keyboard_typing"
+    obs_groups = {
+        "actor": ["policy", "proprio", "perception"],
+        "critic": ["policy", "proprio", "perception"],
+    }
+    actor = SharedEncoderMLPModelCfg(
+        distribution_cfg=RslRlMLPModelCfg.GaussianDistributionCfg(init_std=1.0),
+        obs_normalization=True,
+        hidden_dims=[512, 256, 128],
+        activation="elu",
+        encoder_cfg={"policy": SharedEncoderMLPModelCfg.EncoderCfg(hidden_dims=[256], latent_dim=64)},
+    )
+    critic = SharedEncoderMLPModelCfg(
+        obs_normalization=True,
+        hidden_dims=[512, 256, 128],
+        activation="elu",
+        encoder_cfg={"policy": SharedEncoderMLPModelCfg.EncoderCfg(hidden_dims=[256], latent_dim=64)},
+    )
+    algorithm = RslRlPpoAlgorithmCfg(
+        class_name="isaaclab_tasks.contrib.keyboard.agents.models:SharedEncoderPPO",
+        value_loss_coef=1.0,
+        use_clipped_value_loss=True,
+        clip_param=0.2,
+        entropy_coef=0.005,
+        num_learning_epochs=5,
+        num_mini_batches=4,
+        learning_rate=1.0e-4,
+        schedule="adaptive",
+        gamma=0.995,
+        lam=0.90,
+        desired_kl=0.01,
+        max_grad_norm=1.0,
+    )