Seun/nist gear insertion task example

[seun-doherty]

Summary

• Add a complete NIST gear insertion RL workflow: a Franka Panda robot inserts a medium gear onto a peg on the NIST assembly board using operational-space control and RL Games PPO
• Include step-by-step documentation (environment setup, policy training, evaluation) mirroring the existing Franka lift task pages
• Add RL Games policy wrapper, training script, and YAML config for LSTM-based asymmetric actor-critic

What's Included

Core task (isaaclab_arena/tasks/): task definition, keypoint-squashing rewards, 24-D policy observations with wrist-force feedback, insertion success / gear-drop terminations, domain randomization events

Environment (isaaclab_arena_environments/): OSC action term (asset-relative, EMA smoothing, dead-zones), Franka mimic OSC robot config, environment definition wiring scene + embodiment + task + RL Games

RL infrastructure (isaaclab_arena/policy/, scripts/): RL Games action policy wrapper, training script, PPO hyperparameter config

Documentation (docs/pages/example_workflows/nist_gear_insertion/): 4-page workflow (overview, environment setup, training, evaluation) with GIFs

Asset registry (isaaclab_arena/assets/object_library.py): NIST board, gear, peg, and connector asset definitions

NOTE: Object library paths for added assets will be updated when assets are uploaded to nucleus

[isaaclab-review-bot]

🤖 IsaacLab Review Bot — PR #566

Note: This PR was closed and reopened as #567. Posting review here for the record; the same findings apply to #567.

Overview

Large PR adding a complete NIST gear insertion RL workflow: task definition, keypoint-squashing rewards, 24-D policy observations with wrist-force feedback, OSC action term (asset-relative, EMA smoothing, dead-zones), Franka mimic OSC robot config, RL Games policy wrapper, training script, YAML config, documentation, and asset registry entries. 541 files changed — this includes many changes unrelated to the gear insertion feature (CI, docs restructuring, new embodiments, etc.), which makes isolated review of just the gear insertion feature difficult.

---

Findings

1. 🔴 Critical: _pred_scale is global state — breaks multi-env semantics

File: isaaclab_arena/tasks/rewards/gear_insertion_rewards.py — success_prediction_error.__call__

if true_success.float().mean() >= self._delay_until_ratio:
    self._pred_scale = 1.0

_pred_scale is a scalar that applies to all environments but is flipped to 1.0 based on mean() across all envs. Once flipped, it never resets to 0.0 — not on episode reset, not on environment reset. This means:

• Early episodes where delay_until_ratio hasn't been reached will correctly suppress the penalty
• But once any batch of environments triggers it, all future episodes (even newly reset ones) will be penalized from step 0

Suggestion: Track per-env episode progress or make _pred_scale per-env and reset it in a reset() method.

---

2. 🟡 Warning: Repeated tensor allocation in hot paths

File: isaaclab_arena/tasks/rewards/gear_insertion_rewards.py — _check_gear_position

held_off = (
    torch.tensor(held_gear_base_offset, device=env.device, ...).unsqueeze(0).expand(...)
)
offset = torch.tensor(peg_offset, device=env.device, ...).unsqueeze(0).expand(...)

_check_gear_position is called by gear_insertion_engagement_bonus, gear_insertion_success_bonus, and indirectly by success_prediction_error — all every step. Each call creates new tensors from Python lists. For thousands of envs at 60+ Hz this creates GC pressure.

Suggestion: Pre-compute these tensors once (e.g., in __init__ of the calling classes) and reuse them, similar to how gear_peg_keypoint_squashing already caches self.peg_offset and self.held_gear_base_offset.

---

3. 🟡 Warning: Observation class directly accesses env.action_manager._terms["arm_action"] private API

File: isaaclab_arena/tasks/observations/gear_insertion_observations.py — NistGearInsertionPolicyObservations.__call__

arm_osc_action = env.action_manager._terms["arm_action"]

This accesses the private _terms dict of the action manager. If the action term name changes or the API evolves, this will silently break. The same pattern appears in multiple reward classes in gear_insertion_rewards.py.

Suggestion: Consider adding a public accessor on the action manager or passing the action term via config params rather than reaching into internals. At minimum, add a clear error message if the key isn't found.

---

4. 🟡 Warning: Quaternion canonicalization zeros out z/w components

File: isaaclab_arena/tasks/observations/gear_insertion_observations.py — NistGearInsertionPolicyObservations.__call__

noisy_quat[:, [2, 3]] = 0.0
noisy_quat = noisy_quat * self._flip_quats.unsqueeze(-1)

Zeroing the z and w components of a quaternion after rotation noise application, then multiplying by a random sign, produces a non-unit quaternion that represents a projection rather than a valid rotation. If this is intentional (e.g., only tracking x,y for a 2-DOF orientation), it should be documented. Otherwise, this silently corrupts the rotation information fed to the policy.

---

5. 🟡 Warning: No tests for the new NIST gear insertion task

Despite this being a complete new task workflow, there are no test files specific to the gear insertion task (no test_nist_gear_insertion*.py). The PR adds many other tests (test_assembly_task, test_sorting_task, etc.) but the core new feature — the gear insertion task, its rewards, observations, and OSC action — has no test coverage.

Suggestion: Add at least:

• Unit tests for reward functions (especially the squashing-fn keypoint rewards with known geometry)
• Unit tests for _check_gear_position with controlled inputs
• Integration test that the environment can be created and stepped

---

6. 💡 Nit: body_quat_canonical sign convention may produce discontinuities

File: isaaclab_arena/tasks/observations/gear_insertion_observations.py

sign = torch.where(quat[:, 3:4] < 0, -1.0, 1.0)
return quat * sign

This canonicalizes by flipping so w >= 0. But for orientations near w ≈ 0, small noise flips the sign, creating observation discontinuities. This is a known issue with quaternion canonicalization and may hurt LSTM-based policies. Consider using the make_quat_unique utility from Isaac Lab which handles this more robustly.

---

7. 💡 Nit: gear_peg_height default unused in environment

The NistGearInsertionTask constructor defaults gear_peg_height=0.02, but the environment never passes this parameter — it uses the default. Meanwhile, success_z_fraction=0.20 is explicitly set. The actual success threshold is 0.02 * 0.20 = 0.004m (4mm Z tolerance). This should be documented since it's a critical tuning parameter that's easy to miss.

---

8. 💡 Nit: Duplicate copyright headers

File: isaaclab_arena/scripts/reinforcement_learning/train_rl_games.py

# Copyright (c) 2026, The Isaac Lab Arena Project Developers ...
# Copyright (c) 2025-2026, The Isaac Lab Arena Project Developers.

Two copyright headers at the top — one with 2026 only and one with 2025-2026.

---

Summary

The gear insertion task is well-structured and follows established patterns (Factory-style keypoint rewards, OSC action space, assembly peg-insert benchmarks). The multi-file decomposition (task / observations / rewards / action / environment) is clean and consistent with the existing Arena architecture.

Main concerns: The _pred_scale global state bug (Finding #1) could cause training instabilities. The lack of tests for the core new feature (Finding #5) is a significant gap given the complexity of the reward/observation logic. The hot-path tensor allocations (Finding #2) will hurt training throughput at scale.

The PR's scope is very large (541 files) — consider breaking unrelated changes (CI, docs restructuring, new embodiments) into separate PRs for easier review.