Improve lift rl training example (#529)

## Summary
Update the lift object RL example to have a high success rate model

## Detailed description
- What was the reason for the change?
Existing lift object RL training model success rate is low (~30%), and
the arm motion is unnatural.
- What has been changed?
Add a franka joint control embodiment for RL training to avoid the weird
arm motion from IK version
Update the observation to include joint and target poses only
Fix a bug in the base rls policy so the target pose (task_obs in
addition to policy) is passed to the actor/critic model
Fix a bug resulting ~0 success rate for parallel eval due to incorrect
object/target frame in the success term
Update RL docs with latest models and commands
- What is the impact of this change?
The RL model now gets 70~80% success rate within 1.5h

---------

Co-authored-by: Xinjie Yao <xyao@nvidia.com>

Author: qianl-nv

AGENTS.md

@@ -29,7 +29,7 @@ docker exec isaaclab_arena-latest bash -c "cd /workspaces/isaaclab_arena && \
   --num_steps 10 \
   kitchen_pick_and_place \
   --object cracker_box \
-  --embodiment franka"
+  --embodiment franka_ik"
 ```
 
 ## Common Commands

README.md

@@ -111,7 +111,7 @@ asset_registry = AssetRegistry()
 
 # Select building blocks
 background = asset_registry.get_asset_by_name("kitchen")()
-embodiment = asset_registry.get_asset_by_name("franka")()
+embodiment = asset_registry.get_asset_by_name("franka_ik")()
 cracker_box = asset_registry.get_asset_by_name("cracker_box")()
 tomato_soup_can = asset_registry.get_asset_by_name("tomato_soup_can")()
 

docs/images/lift_object_rl_parallel.gif

@@ -100,7 +100,7 @@ The following code snippet shows a simple example(pick up a tomato soup can and
 
 .. code-block:: python
 
-   embodiment = asset_registry.get_asset_by_name("franka")(enable_cameras=True)
+   embodiment = asset_registry.get_asset_by_name("franka_ik")(enable_cameras=True)
    background = asset_registry.get_asset_by_name("kitchen")()
    tomato_soup_can = asset_registry.get_asset_by_name("tomato_soup_can")()
    destination_location = ObjectReference(

docs/images/lift_object_rl_task.gif

@@ -119,7 +119,7 @@ Environment Integration
 .. code-block:: python
 
    # Embodiment creation with camera support
-   embodiment = asset_registry.get_asset_by_name("franka")(
+   embodiment = asset_registry.get_asset_by_name("franka_ik")(
        enable_cameras=True
    )
 
@@ -144,7 +144,7 @@ Usage Examples
 
 .. code-block:: python
 
-   franka = asset_registry.get_asset_by_name("franka")(enable_cameras=True)
+   franka = asset_registry.get_asset_by_name("franka_ik")(enable_cameras=True)
    task = PickAndPlaceTask(pick_object, destination, background)
 
 **Humanoid Control Modes**

docs/index.rst

@@ -33,7 +33,7 @@ Creating an Environment Example
 .. code-block:: python
 
    # Component creation
-   embodiment = asset_registry.get_asset_by_name("franka")()
+   embodiment = asset_registry.get_asset_by_name("franka_ik")()
    background = asset_registry.get_asset_by_name("kitchen")()
    pick_object = asset_registry.get_asset_by_name("cracker_box")()
    pick_object.set_initial_pose(Pose(position_xyz=(0.4, 0.0, 0.1)))

docs/pages/concepts/concept_embodiment_design.rst

@@ -18,8 +18,8 @@ builds the environment, and registers it with gym so IsaacLab's script can find
      --external_callback isaaclab_arena.environments.isaaclab_interop.environment_registration_callback \
      --task lift_object \
      --rl_training_mode \
-     --num_envs 512 \
-     --max_iterations 12000
+     --num_envs 4096 \
+     --max_iterations 2000
 
 .. tip::
 
@@ -55,8 +55,8 @@ For example, to train with relu activation and a higher learning rate:
      --external_callback isaaclab_arena.environments.isaaclab_interop.environment_registration_callback \
      --task lift_object \
      --rl_training_mode \
-     --num_envs 512 \
-     --max_iterations 12000 \
+     --num_envs 4096 \
+     --max_iterations 2000 \
      agent.policy.activation=relu \
      agent.algorithm.learning_rate=0.001
 
@@ -112,15 +112,15 @@ Add ``--distributed`` to spread environments across all available GPUs:
      --external_callback isaaclab_arena.environments.isaaclab_interop.environment_registration_callback \
      --task lift_object \
      --rl_training_mode \
-     --num_envs 512 \
-     --max_iterations 12000 \
+     --num_envs 4096\
+     --max_iterations 2000 \
      --distributed
 
 
 Expected Results
 ^^^^^^^^^^^^^^^^
 
-After 12,000 iterations (~6 hours on a single GPU with 512 environments), the trained
+After 2,000 iterations (~40 minutes on a single GPU with 4096 environments), the trained
 policy should reliably grasp and lift objects to commanded target positions.
 
 .. image:: ../../../images/lift_object_rl_task.gif

docs/pages/concepts/concept_environment_design.rst

@@ -21,13 +21,12 @@ or you can download a pre-trained one as described below.
    .. code-block:: bash
 
       hf download \
-         nvidia/IsaacLab-Arena-Lift-Object-RL \
-         model_11999.pt \
+         nvidia/Arena-Franka-Lift-Object-RL-Task \
          --local-dir $MODELS_DIR/lift_object_checkpoint
 
    After downloading, the checkpoint is at:
 
-   ``$MODELS_DIR/lift_object_checkpoint/model_11999.pt``
+   ``$MODELS_DIR/lift_object_checkpoint/model_1999.pt``
 
    Replace checkpoint paths in the examples below with this path.
 
@@ -50,14 +49,14 @@ Method 1: Single Environment Evaluation
    python isaaclab_arena/evaluation/policy_runner.py \
      --visualizer kit \
      --policy_type rsl_rl \
-     --num_steps 1000 \
-     --checkpoint_path logs/rsl_rl/generic_experiment/2026-01-28_17-26-10/model_11999.pt \
+     --num_episodes 20 \
+     --checkpoint_path $MODELS_DIR/lift_object_checkpoint/model_1999.pt \
      lift_object
 
 .. note::
 
-   If you downloaded the pre-trained model from Hugging Face, replace the checkpoint path with:
-   ``$MODELS_DIR/lift_object_checkpoint/model_11999.pt``
+   If you train the model yourself, the checkpoint path is typically in the ``logs/rsl_rl/generic_experiment/`` directory.
+   Replace the checkpoint path with the path to your own checkpoint.
 
 Policy-specific arguments (``--policy_type``, ``--checkpoint_path``, etc.) must come **before** the
 environment name. Environment-specific arguments (``--object``, ``--embodiment``, etc.) must come
@@ -67,7 +66,7 @@ At the end of the run, metrics are printed to the console:
 
 .. code-block:: text
 
-   Metrics: {'success_rate': 0.85, 'num_episodes': 12}
+   Metrics: {'success_rate': 0.81, 'num_episodes': 12}
 
 
 Method 2: Parallel Environment Evaluation
@@ -79,21 +78,28 @@ For more statistically significant results, run across many environments in para
 
    python isaaclab_arena/evaluation/policy_runner.py \
      --policy_type rsl_rl \
-     --num_steps 5000 \
+     --num_episodes 1024 \
      --num_envs 64 \
-     --checkpoint_path logs/rsl_rl/generic_experiment/2026-01-28_17-26-10/model_11999.pt \
-     --headless \
+     --env_spacing 2.5 \
+     --visualizer kit \
+     --checkpoint_path $MODELS_DIR/lift_object_checkpoint/model_1999.pt \
      lift_object
 
 .. code-block:: text
 
-   Metrics: {'success_rate': 0.83, 'num_episodes': 156}
+   Metrics: {'success_rate': 0.72, 'num_episodes': 1024}
+
+.. image:: ../../../images/lift_object_rl_parallel.gif
+   :align: center
+   :height: 400px
 
 
 Method 3: Batch Evaluation
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^
 
 To evaluate multiple checkpoints in sequence, use ``eval_runner.py`` with a JSON config.
+Here we evaluate the models you trained yourself.
+The checkpoint path should be replaced with the timestamp of your training run in the ``logs/rsl_rl/generic_experiment/`` directory.
 
 **1. Create an evaluation config**
 
@@ -102,20 +108,30 @@ Create a file ``eval_config.json``:
 .. code-block:: json
 
    {
-     "policy_runner_args": {
-       "policy_type": "rsl_rl",
-       "num_steps": 5000,
-       "num_envs": 64,
-       "headless": true
-     },
-     "evaluations": [
+     "jobs": [
        {
-         "checkpoint_path": "logs/rsl_rl/generic_experiment/2026-01-28_17-26-10/model_5999.pt",
-         "environment": "lift_object"
+         "name": "lift_object_model_1000",
+         "policy_type": "rsl_rl",
+         "num_episodes": 1024,
+         "arena_env_args": {
+           "environment": "lift_object",
+           "num_envs": 64
+         },
+         "policy_config_dict": {
+           "checkpoint_path": "logs/rsl_rl/generic_experiment/<timestamp>/model_1000.pt"
+         }
        },
        {
-         "checkpoint_path": "logs/rsl_rl/generic_experiment/2026-01-28_17-26-10/model_11999.pt",
-         "environment": "lift_object"
+         "name": "lift_object_model_1999",
+         "policy_type": "rsl_rl",
+         "num_episodes": 1024,
+         "arena_env_args": {
+           "environment": "lift_object",
+           "num_envs": 64
+         },
+         "policy_config_dict": {
+           "checkpoint_path": "logs/rsl_rl/generic_experiment/<timestamp>/model_1999.pt"
+         }
        }
      ]
    }
@@ -128,16 +144,18 @@ Create a file ``eval_config.json``:
 
 .. code-block:: text
 
-   Evaluating checkpoint 1/2: model_5999.pt
-   Metrics: {'success_rate': 0.72, 'num_episodes': 152}
+   ======================================================================
+   METRICS SUMMARY
+   ======================================================================
 
-   Evaluating checkpoint 2/2: model_11999.pt
-   Metrics: {'success_rate': 0.85, 'num_episodes': 156}
+   lift_object_model_1000:
+   num_episodes                         1024
+   success_rate                       0.6526
 
-   Summary:
-   ========================================
-   model_5999.pt  | Success: 72% | Episodes: 152
-   model_11999.pt | Success: 85% | Episodes: 156
+   lift_object_model_1999:
+   num_episodes                         1024
+   success_rate                       0.7408
+   ======================================================================
 
 
 Understanding the Metrics

docs/pages/example_workflows/reinforcement_learning/step_2_policy_training.rst

@@ -75,7 +75,7 @@ def get_pipeline_builder(self, embodiment: object) -> Callable:
 @register_retargeter
 class FrankaKeyboardRetargeter(RetargetterBase):
     device = "keyboard"
-    embodiment = "franka"
+    embodiment = "franka_ik"
 
     def __init__(self):
         pass
@@ -87,7 +87,7 @@ def get_pipeline_builder(self, embodiment: object) -> Callable | None:
 @register_retargeter
 class FrankaSpaceMouseRetargeter(RetargetterBase):
     device = "spacemouse"
-    embodiment = "franka"
+    embodiment = "franka_ik"
 
     def __init__(self):
         pass