Fixes lift environment for Franka robot (#264)

# Description

Adds the Franka lift RL environment. It also checks that the reach
environment is working.

## Type of change

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

## Checklist

- [x] I have run the [`pre-commit` checks](https://pre-commit.com/) with
`./orbit.sh --format`
- [ ] I have made corresponding changes to the documentation
- [x] My changes generate no new warnings
- [ ] I have added tests that prove my fix is effective or that my
feature works
- [ ] I have updated the changelog and the corresponding version in the
extension's `config/extension.toml` file
- [x] I have added my name to the `CONTRIBUTORS.md` or my name already
exists there

source/extensions/omni.isaac.orbit/omni/isaac/orbit/assets/config/franka.py

@@ -52,22 +52,22 @@
             joint_names_expr=["panda_joint[1-4]"],
             effort_limit=87.0,
             velocity_limit=100.0,
-            stiffness=800.0,
-            damping=40.0,
+            stiffness=80.0,
+            damping=4.0,
         ),
         "panda_forearm": ImplicitActuatorCfg(
             joint_names_expr=["panda_joint[5-7]"],
-            effort_limit=12.0,
+            effort_limit=87.0,
             velocity_limit=100.0,
-            stiffness=800.0,
-            damping=40.0,
+            stiffness=80.0,
+            damping=4.0,
         ),
         "panda_hand": ImplicitActuatorCfg(
             joint_names_expr=["panda_finger_joint.*"],
             effort_limit=200.0,
             velocity_limit=0.2,
-            stiffness=1e5,
-            damping=1e3,
+            stiffness=2e3,
+            damping=1e2,
         ),
     },
     soft_joint_pos_limit_factor=1.0,

source/extensions/omni.isaac.orbit/omni/isaac/orbit/utils/math.py

@@ -3,10 +3,7 @@
 #
 # SPDX-License-Identifier: BSD-3-Clause
 
-"""Sub-module containing utilities for various math operations.
-
-Some of these are imported from the module `omni.isaac.core.utils.torch` for convenience.
-"""
+"""Sub-module containing utilities for various math operations."""
 
 from __future__ import annotations
 
@@ -690,7 +687,7 @@ def is_identity_pose(pos: torch.tensor, rot: torch.tensor) -> bool:
 
 # @torch.jit.script
 def combine_frame_transforms(
-    t01: torch.Tensor, q01: torch.Tensor, t12: torch.Tensor = None, q12: torch.Tensor = None
+    t01: torch.Tensor, q01: torch.Tensor, t12: torch.Tensor | None = None, q12: torch.Tensor | None = None
 ) -> tuple[torch.Tensor, torch.Tensor]:
     r"""Combine transformations between two reference frames into a stationary frame.
 
@@ -701,7 +698,9 @@ def combine_frame_transforms(
         t01: Position of frame 1 w.r.t. frame 0. Shape is (N, 3).
         q01: Quaternion orientation of frame 1 w.r.t. frame 0 in (w, x, y, z). Shape is (N, 4).
         t12: Position of frame 2 w.r.t. frame 1. Shape is (N, 3).
+            Defaults to None, in which case the position is assumed to be zero.
         q12: Quaternion orientation of frame 2 w.r.t. frame 1 in (w, x, y, z). Shape is (N, 4).
+            Defaults to None, in which case the orientation is assumed to be identity.
 
     Returns:
         A tuple containing the position and orientation of frame 2 w.r.t. frame 0.
@@ -721,9 +720,9 @@ def combine_frame_transforms(
     return t02, q02
 
 
-@torch.jit.script
+# @torch.jit.script
 def subtract_frame_transforms(
-    t01: torch.Tensor, q01: torch.Tensor, t02: torch.Tensor, q02: torch.Tensor
+    t01: torch.Tensor, q01: torch.Tensor, t02: torch.Tensor | None = None, q02: torch.Tensor | None = None
 ) -> tuple[torch.Tensor, torch.Tensor]:
     r"""Subtract transformations between two reference frames into a stationary frame.
 
@@ -734,18 +733,25 @@ def subtract_frame_transforms(
         t01: Position of frame 1 w.r.t. frame 0. Shape is (N, 3).
         q01: Quaternion orientation of frame 1 w.r.t. frame 0 in (w, x, y, z). Shape is (N, 4).
         t02: Position of frame 2 w.r.t. frame 0. Shape is (N, 3).
+            Defaults to None, in which case the position is assumed to be zero.
         q02: Quaternion orientation of frame 2 w.r.t. frame 0 in (w, x, y, z). Shape is (N, 4).
+            Defaults to None, in which case the orientation is assumed to be identity.
 
     Returns:
         A tuple containing the position and orientation of frame 2 w.r.t. frame 1.
         Shape of the tensors are (N, 3) and (N, 4) respectively.
     """
     # compute orientation
     q10 = quat_inv(q01)
-    q12 = quat_mul(q10, q02)
+    if q02 is not None:
+        q12 = quat_mul(q10, q02)
+    else:
+        q12 = q10
     # compute translation
-    t12 = quat_apply(q10, t02 - t01)
-
+    if t02 is not None:
+        t12 = quat_apply(q10, t02 - t01)
+    else:
+        t12 = quat_apply(q10, -t01)
     return t12, q12
 
 
@@ -816,7 +822,7 @@ def apply_delta_pose(
         source_pos: Position of source frame. Shape is (N, 3).
         source_rot: Quaternion orientation of source frame in (w, x, y, z). Shape is (N, 4)..
         delta_pose: Position and orientation displacements. Shape is (N, 6).
-        eps: The tolerance to consider orientation displacement as zero.
+        eps: The tolerance to consider orientation displacement as zero. Defaults to 1.0e-6.
 
     Returns:
         A tuple containing the displaced position and orientation frames.
@@ -864,7 +870,9 @@ def transform_points(
     Args:
         points: Points to transform. Shape is (N, P, 3) or (P, 3).
         pos: Position of the target frame. Shape is (N, 3) or (3,).
+            Defaults to None, in which case the position is assumed to be zero.
         quat: Quaternion orientation of the target frame in (w, x, y, z). Shape is (N, 4) or (4,).
+            Defaults to None, in which case the orientation is assumed to be identity.
 
     Returns:
         Transformed points in the target frame. Shape is (N, P, 3) or (P, 3).

source/extensions/omni.isaac.orbit_tasks/omni/isaac/orbit_tasks/manipulation/lift/__init__.py

@@ -3,29 +3,7 @@
 #
 # SPDX-License-Identifier: BSD-3-Clause
 
-"""
-Environment for lifting an object with fixed-base robot.
-"""
+"""Configurations for the object lift environments."""
 
-import gymnasium as gym
-
-from . import agents
-
-##
-# Register Gym environments.
-##
-
-gym.register(
-    id="Isaac-Lift-Franka-v0",
-    entry_point="omni.isaac.orbit.envs:RLTaskEnv",
-    disable_env_checker=True,
-    kwargs={
-        "env_cfg_entry_point": f"{__name__}.lift_env_cfg:LiftEnvCfg",
-        "rl_games_cfg_entry_point": f"{agents.__name__}:rl_games_ppo_cfg.yaml",
-        "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:LIFT_RSL_RL_PPO_CFG",
-        "skrl_cfg_entry_point": f"{agents.__name__}:skrl_ppo_cfg.yaml",
-        "sb3_cfg_entry_point": f"{agents.__name__}:sb3_ppo_cfg.yaml",
-        "robomimic_bc_cfg_entry_point": f"{agents.__name__}.robomimic:bc.json",
-        "robomimic_bcq_cfg_entry_point": f"{agents.__name__}.robomimic:bcq.json",
-    },
-)
+# We leave this file empty since we don't want to expose any configs in this package directly.
+# We still need this file to import the "config" module in the parent package.

source/extensions/omni.isaac.orbit_tasks/omni/isaac/orbit_tasks/manipulation/lift/config/__init__.py

@@ -0,0 +1,9 @@
+# Copyright (c) 2022-2023, The ORBIT Project Developers.
+# All rights reserved.
+#
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""Configurations for the object lift environments."""
+
+# We leave this file empty since we don't want to expose any configs in this package directly.
+# We still need this file to import the "config" module in the parent package.

source/extensions/omni.isaac.orbit_tasks/omni/isaac/orbit_tasks/manipulation/lift/config/franka/__init__.py

@@ -0,0 +1,32 @@
+# Copyright (c) 2022-2023, The ORBIT Project Developers.
+# All rights reserved.
+#
+# SPDX-License-Identifier: BSD-3-Clause
+
+
+import gymnasium as gym
+
+from . import agents, env_cfg
+
+##
+# Register Gym environments.
+##
+gym.register(
+    id="Isaac-Lift-Cube-Franka-v0",
+    entry_point="omni.isaac.orbit.envs:RLTaskEnv",
+    kwargs={
+        "env_cfg_entry_point": env_cfg.FrankaCubeLiftEnvCfg,
+        "rsl_rl_cfg_entry_point": agents.rsl_rl_cfg.LiftCubePPORunnerCfg,
+    },
+    disable_env_checker=True,
+)
+
+gym.register(
+    id="Isaac-Lift-Cube-Franka-Play-v0",
+    entry_point="omni.isaac.orbit.envs:RLTaskEnv",
+    kwargs={
+        "env_cfg_entry_point": env_cfg.FrankaCubeLiftEnvCfg_PLAY,
+        "rsl_rl_cfg_entry_point": agents.rsl_rl_cfg.LiftCubePPORunnerCfg,
+    },
+    disable_env_checker=True,
+)

source/extensions/omni.isaac.orbit_tasks/omni/isaac/orbit_tasks/manipulation/lift/config/franka/agents/__init__.py

@@ -0,0 +1,6 @@
+# Copyright (c) 2022-2023, The ORBIT Project Developers.
+# All rights reserved.
+#
+# SPDX-License-Identifier: BSD-3-Clause
+
+from . import rsl_rl_cfg  # noqa: F401, F403

source/extensions/omni.isaac.orbit_tasks/omni/isaac/orbit_tasks/manipulation/lift/agents/rsl_rl_ppo_cfg.py → source/extensions/omni.isaac.orbit_tasks/omni/isaac/orbit_tasks/manipulation/lift/config/franka/agents/rsl_rl_cfg.py

@@ -3,40 +3,39 @@
 #
 # SPDX-License-Identifier: BSD-3-Clause
 
+from omni.isaac.orbit.utils import configclass
+
 from omni.isaac.orbit_tasks.utils.wrappers.rsl_rl import (
     RslRlOnPolicyRunnerCfg,
     RslRlPpoActorCriticCfg,
     RslRlPpoAlgorithmCfg,
 )
 
-LIFT_RSL_RL_PPO_CFG = RslRlOnPolicyRunnerCfg(
-    num_steps_per_env=96,
-    max_iterations=700,
-    save_interval=50,
-    experiment_name="lift",
-    run_name="",
-    resume=False,
-    load_run=-1,
-    load_checkpoint=-1,
-    empirical_normalization=False,
-    policy=RslRlPpoActorCriticCfg(
-        init_noise_std=1.0,
+
+@configclass
+class LiftCubePPORunnerCfg(RslRlOnPolicyRunnerCfg):
+    num_steps_per_env = 24
+    max_iterations = 1500
+    save_interval = 50
+    experiment_name = "franka_lift"
+    empirical_normalization = False
+    policy = RslRlPpoActorCriticCfg(
+        init_noise_std=0.8,
         actor_hidden_dims=[256, 128, 64],
         critic_hidden_dims=[256, 128, 64],
         activation="elu",
-    ),
-    algorithm=RslRlPpoAlgorithmCfg(
+    )
+    algorithm = RslRlPpoAlgorithmCfg(
         value_loss_coef=1.0,
         use_clipped_value_loss=True,
         clip_param=0.2,
-        entropy_coef=0.01,
+        entropy_coef=0.006,
         num_learning_epochs=5,
         num_mini_batches=4,
-        learning_rate=1.0e-3,
+        learning_rate=1.0e-4,
         schedule="adaptive",
-        gamma=0.99,
+        gamma=0.98,
         lam=0.95,
         desired_kl=0.01,
         max_grad_norm=1.0,
-    ),
-)
+    )

source/extensions/omni.isaac.orbit_tasks/omni/isaac/orbit_tasks/manipulation/lift/config/franka/env_cfg.py

@@ -0,0 +1,94 @@
+# Copyright (c) 2022-2023, The ORBIT Project Developers.
+# All rights reserved.
+#
+# SPDX-License-Identifier: BSD-3-Clause
+
+from omni.isaac.orbit.assets import RigidObjectCfg
+from omni.isaac.orbit.sensors import FrameTransformerCfg
+from omni.isaac.orbit.sensors.frame_transformer.frame_transformer_cfg import OffsetCfg
+from omni.isaac.orbit.sim.schemas.schemas_cfg import RigidBodyPropertiesCfg
+from omni.isaac.orbit.sim.spawners.from_files.from_files_cfg import UsdFileCfg
+from omni.isaac.orbit.utils import configclass
+from omni.isaac.orbit.utils.assets import ISAAC_NUCLEUS_DIR
+
+from omni.isaac.orbit_tasks.manipulation.lift import mdp
+from omni.isaac.orbit_tasks.manipulation.lift.lift_env_cfg import LiftEnvCfg
+
+##
+# Pre-defined configs
+##
+# isort: off
+from omni.isaac.orbit.assets.config.franka import FRANKA_PANDA_ARM_WITH_PANDA_HAND_CFG
+from omni.isaac.orbit.markers.config import FRAME_MARKER_CFG
+
+
+@configclass
+class FrankaCubeLiftEnvCfg(LiftEnvCfg):
+    def __post_init__(self):
+        # post init of parent
+        super().__post_init__()
+
+        # Set Franka as robot
+        self.scene.robot = FRANKA_PANDA_ARM_WITH_PANDA_HAND_CFG.replace(prim_path="{ENV_REGEX_NS}/Robot")
+
+        # Set actions for the specific robot type (franka)
+        self.actions.body_joint_pos = mdp.JointPositionActionCfg(
+            asset_name="robot", joint_names=["panda_joint.*"], scale=1.0, use_default_offset=True
+        )
+        self.actions.finger_joint_pos = mdp.BinaryJointPositionActionCfg(
+            asset_name="robot",
+            joint_names=["panda_finger.*"],
+            open_command_expr={"panda_finger_.*": 0.3},
+            close_command_expr={"panda_finger_.*": 0.0},
+        )
+        # Set the body name for the end effector
+        self.commands.body_name = "panda_hand"
+
+        # Set Cube as object
+        self.scene.object = RigidObjectCfg(
+            prim_path="{ENV_REGEX_NS}/Object",
+            init_state=RigidObjectCfg.InitialStateCfg(pos=[0.5, 0, 0.055], rot=[1, 0, 0, 0]),
+            spawn=UsdFileCfg(
+                usd_path=f"{ISAAC_NUCLEUS_DIR}/Props/Blocks/DexCube/dex_cube_instanceable.usd",
+                scale=(0.8, 0.8, 0.8),
+                rigid_props=RigidBodyPropertiesCfg(
+                    solver_position_iteration_count=16,
+                    solver_velocity_iteration_count=1,
+                    max_angular_velocity=1000.0,
+                    max_linear_velocity=1000.0,
+                    max_depenetration_velocity=5.0,
+                    disable_gravity=False,
+                ),
+            ),
+        )
+
+        # Listens to the required transforms
+        marker_cfg = FRAME_MARKER_CFG.copy()
+        marker_cfg.markers["frame"].scale = (0.1, 0.1, 0.1)
+        marker_cfg.prim_path = "/Visuals/FrameTransformer"
+        self.scene.ee_frame = FrameTransformerCfg(
+            prim_path="{ENV_REGEX_NS}/Robot/panda_link0",
+            debug_vis=False,
+            visualizer_cfg=marker_cfg,
+            target_frames=[
+                FrameTransformerCfg.FrameCfg(
+                    prim_path="{ENV_REGEX_NS}/Robot/panda_hand",
+                    name="end_effector",
+                    offset=OffsetCfg(
+                        pos=[0.0, 0.0, 0.1034],
+                    ),
+                ),
+            ],
+        )
+
+
+@configclass
+class FrankaCubeLiftEnvCfg_PLAY(FrankaCubeLiftEnvCfg):
+    def __post_init__(self):
+        # post init of parent
+        super().__post_init__()
+        # make a smaller scene for play
+        self.scene.num_envs = 50
+        self.scene.env_spacing = 2.5
+        # disable randomization for play
+        self.observations.policy.enable_corruption = False