Probelms with physics

[Icon45]

Hi,

i'm trying to get some boxes to behave physically correct in an Isaac Sim environment (4.5.0) that is being launched and used by Isaac Lab (2.1.0).

Here is my code (as i'm quit new to Isaac feel free to critizise my code or point out other issues that you come across):

# ───────────────────────── 0)  CLI & SimulationApp ─────────────────────────
import argparse, math, pathlib, random
from typing import Optional, Sequence, List, Tuple

import numpy as np
import torch
import torch.nn as nn

from isaaclab.app import AppLauncher

parser = argparse.ArgumentParser(description="HC10 PPO Training (skrl Edition)")
parser.add_argument("--num_envs", type=int, default=80, help="Parallele Sub-Envs")
AppLauncher.add_app_launcher_args(parser)      # --headless, --width, ...
args_cli = parser.parse_args()

# →  SimulationApp/Kit starten  (muss VOR allen Omniverse-Imports passieren)
app_launcher   = AppLauncher(args_cli)
simulation_app = app_launcher.app           # noqa: F841

# ───────────────────────── 1)  Hyper- & Train-Parameter ────────────────────
TOTAL_TIMESTEPS = 2_000_000
ROLLOUT_LENGTH  = 4096                      # Schritte pro Env bis Update
BATCH_SIZE      = 1024                      # Mini-Batch-Größe (skrl = mini_batches)
LEARNING_RATE   = 2.5e-4
SEED: Optional[int] = 42 
DOF = 6                                     # HC10 hat 6 Achsen

# ───────────────────────── 2)  Isaac-Lab Imports  ──────────────────────────
import isaaclab.sim as sim_utils
from isaaclab.scene import InteractiveSceneCfg
from isaaclab.envs import DirectRLEnv, DirectRLEnvCfg
from isaaclab.assets.articulation import Articulation, ArticulationCfg
from isaaclab.actuators import ImplicitActuatorCfg
from isaaclab.sim.spawners.from_files import UsdFileCfg, spawn_from_usd
from isaaclab.utils import configclass
from isaaclab.utils.math import quat_from_euler_xyz as q_from_euler
from pxr import Usd, UsdPhysics, PhysxSchema
from isaacsim.core.utils import prims as prim_utils  # optional
from gymnasium import spaces

# ───────────────────────── 2b)  PhysX-State-API für harte Resets ──────────
# Isaac Lab 2.1: acquire_dynamic_control_interface() liegt im Untermodul _dynamic_control
from omni.isaac.dynamic_control import _dynamic_control as dc
import carb
import omni.kit.app  # optional: Extension sicher aktivieren

# ───────────────────────── 3)  s k r l Imports  ────────────────────────────
from skrl.envs.wrappers.torch import wrap_env
from skrl.memories.torch.random import RandomMemory
from skrl.models.torch import GaussianMixin, DeterministicMixin, Model
from skrl.agents.torch.ppo import PPO, PPO_DEFAULT_CONFIG
from skrl.trainers.torch import SequentialTrainer

# ───────────────────────── 4)  Pfade & Konstanten  ─────────────────────────
USD_ALL    = "description/yhc10Pneumatik/configuration/yhc10Pneumatik_base.usd"

MOUNT_PATH = "/World/envs/env_0/sortieranlage_full"
REL_ROOT   = "/sortieranlage/world"
ROBOT_ROOT = MOUNT_PATH + REL_ROOT

BLUE, ORANGE = (0.2, 0.4, 1.0), (1.0, 0.5, 0.2)
SMALL_POS: Sequence[tuple[float, float]] = (
    (-0.11, 0.90), (0.00, 0.90), (0.11, 0.90),
    (-0.11, 0.60), (0.00, 0.60), (0.11, 0.60),
    (-0.25, 0.69), (-0.25, 0.81), (0.25, 0.69), (0.25, 0.81),
)

BOX_BIG_CFG = sim_utils.CuboidCfg(
    size=(0.1, 0.2, 0.1),
    rigid_props=sim_utils.RigidBodyPropertiesCfg(kinematic_enabled=False, disable_gravity=False),
    mass_props=sim_utils.MassPropertiesCfg(mass=2.0),
    collision_props=sim_utils.CollisionPropertiesCfg(),
    visual_material=sim_utils.PreviewSurfaceCfg(diffuse_color=BLUE),
)
BOX_SMALL_CFG = sim_utils.CuboidCfg(
    size=(0.1, 0.1, 0.1),
    rigid_props=sim_utils.RigidBodyPropertiesCfg(kinematic_enabled=False, disable_gravity=False),
    mass_props=sim_utils.MassPropertiesCfg(mass=1.0),
    collision_props=sim_utils.CollisionPropertiesCfg(),
    visual_material=sim_utils.PreviewSurfaceCfg(diffuse_color=ORANGE),
)

# ───────────────────────── 5)  Task-Klasse  ───────────────────────────────
class ClearanceTask:
    """Reset, Step, Reward – Multi-Env-fähig durch batched Writes"""

    _ENV_ROOT = "/World/envs/env_{eid}"

    def __init__(self, scene, robot, dof: int):
        self.scene, self.robot, self._dof = scene, robot, dof
        self.num_small = len(SMALL_POS)

        self.big_xy: dict[int, np.ndarray] = {}
        self.small_xy: dict[int, np.ndarray] = {}
        self.big_box_path: dict[int, str] = {}
        self.small_box_paths: dict[int, List[str]] = {}

        # PhysX-State-API holen
        self._dc = dc.acquire_dynamic_control_interface()
        if not self._dc:
            raise RuntimeError("dynamic_control interface not available")

    # ------------------------------------------------------------------
    # Helper
    # ------------------------------------------------------------------
    @staticmethod
    def _quat_z(yaw: float, dev: torch.device) -> torch.Tensor:
        return q_from_euler(torch.tensor([0.0], device=dev),
                            torch.tensor([0.0], device=dev),
                            torch.tensor([yaw],  device=dev)).squeeze(0)

    @staticmethod
    def _wxyz_to_xyzw(q_wxyz: torch.Tensor) -> Tuple[float, float, float, float]:
        q = q_wxyz.detach().float().cpu().tolist()
        return (float(q[1]), float(q[2]), float(q[3]), float(q[0]))

    # ------------------------------------------------------------------
    # USD / PhysX Utilities
    # ------------------------------------------------------------------
    def _has_rigid_api(self, path: str) -> bool:
        prim = self.scene.stage.GetPrimAtPath(path)
        if not prim or not prim.IsValid():
            return False
        return prim.HasAPI(UsdPhysics.RigidBodyAPI) or prim.HasAPI(PhysxSchema.PhysxRigidBodyAPI)

    def _find_child_rigid_body(self, root_path: str):
        if self._has_rigid_api(root_path):
            rb = self._dc.get_rigid_body(root_path)
            if rb:
                return rb, root_path

        root_prim = self.scene.stage.GetPrimAtPath(root_path)
        if not root_prim or not root_prim.IsValid():
            return None, root_path

        for prim in Usd.PrimRange(root_prim):
            p = prim.GetPath().pathString
            if self._has_rigid_api(p):
                rb = self._dc.get_rigid_body(p)
                if rb:
                    return rb, p
        return None, root_path

    def _ensure_dynamic_rigid_body(self, prim_path: str, mass: float):
        """Wendet USD- und PhysX-APIs an und stellt dynamische Eigenschaften ein."""
        stage = self.scene.stage
        prim = stage.GetPrimAtPath(prim_path)
        if not prim or not prim.IsValid():
            return

        # 1) USD RigidBody
        if not prim.HasAPI(UsdPhysics.RigidBodyAPI):
            UsdPhysics.RigidBodyAPI.Apply(prim)

        # 2) PhysX RigidBody
        if not prim.HasAPI(PhysxSchema.PhysxRigidBodyAPI):
            PhysxSchema.PhysxRigidBodyAPI.Apply(prim)
        physx_rb = PhysxSchema.PhysxRigidBodyAPI(prim)

        try:
            physx_rb.GetKinematicEnabledAttr().Set(False)
        except AttributeError:
            try:
                physx_rb.CreateKinematicEnabledAttr(False)
            except AttributeError:
                pass

        try:
            physx_rb.GetDisableGravityAttr().Set(False)
        except AttributeError:
            try:
                physx_rb.CreateDisableGravityAttr(False)
            except AttributeError:
                pass

        try:                                                 #  <-- NEU
            physx_rb.GetDisableSleepingAttr().Set(True)      #  <-- NEU
        except AttributeError:                               #  <-- NEU
            try:                                             #  <-- NEU
                physx_rb.CreateDisableSleepingAttr(True)     #  <-- NEU
            except AttributeError:                           #  <-- NEU
                pass                                         #  <-- NEU

        # 3) Masse
        if not prim.HasAPI(UsdPhysics.MassAPI):
            UsdPhysics.MassAPI.Apply(prim)
        mass_api = UsdPhysics.MassAPI(prim)
        if not mass_api.GetMassAttr().HasAuthoredValueOpinion():
            mass_api.CreateMassAttr(float(mass))

        # 4) CollisionAPI (optional)
        if not prim.HasAPI(UsdPhysics.CollisionAPI):
            UsdPhysics.CollisionAPI.Apply(prim)

    def _set_rb_state(self, prim_path: str, p_xyz, q_xyzw):
        """Rigid-Body teleportieren und **aufwecken**."""
        rb, _ = self._find_child_rigid_body(prim_path)
        if not rb:
            return

        st = self._dc.get_rigid_body_state(rb)
        st.pose.p = carb.Float3(float(p_xyz[0]), float(p_xyz[1]), float(p_xyz[2]))
        st.pose.r = carb.Float4(float(q_xyzw[0]), float(q_xyzw[1]), float(q_xyzw[2]), float(q_xyzw[3]))
        st.linvel = carb.Float3(0.0, 0.0, 0.0)
        st.angvel = carb.Float3(0.0, 0.0, 0.0)
        if hasattr(st, "sleeping"):
            st.sleeping = False
        self._dc.set_rigid_body_state(rb, st, True)          # wake_up=True

    # ------------------------------------------------------------------
    # Public API
    # ------------------------------------------------------------------
    def reset(self, eid: int):
        dev = self.robot.data.joint_pos.device
        self.reset_batch(torch.tensor([eid], dtype=torch.int32, device=dev))

    def reset_batch(self, env_ids: torch.Tensor):
        dev = self.robot.data.joint_pos.device
        env_ids = env_ids.to(device=dev, dtype=torch.int32)
        batch_size = int(env_ids.shape[0])

        pos = torch.zeros((batch_size, self._dof), device=dev)
        vel = torch.zeros((batch_size, self._dof), device=dev)
        self.robot.write_joint_state_to_sim(pos, vel, env_ids=env_ids)

        for eid in env_ids.tolist():
            env_root = self._ENV_ROOT.format(eid=eid)
            self.scene.stage.DefinePrim(env_root, "Xform")

            if eid not in self.big_box_path:
                self._spawn_boxes(eid, env_root, dev)
            else:
                self._teleport_boxes(eid, dev)

            self.big_xy[eid] = np.array([0.0, 0.75], dtype=np.float32)
            self.small_xy[eid] = np.asarray(SMALL_POS, dtype=np.float32)

    # ------------------------------------------------------------------
    # Spawn / Teleport Helper
    # ------------------------------------------------------------------
    def _spawn_boxes(self, eid: int, env_root: str, dev: torch.device):
        global BOX_BIG_CFG, BOX_SMALL_CFG

        # ─── große Box ─────────────────────────────────────────
        big_path = f"{env_root}/box_big"
        BOX_BIG_CFG.func(
            big_path, BOX_BIG_CFG,
            translation=(0.0, 3.75, 1.0),
            orientation=self._quat_z(math.pi / 2, dev)
        )
        self._ensure_dynamic_rigid_body(big_path, mass=2.0)
        self.big_box_path[eid] = big_path

        # ─── kleine Boxen ─────────────────────────────────────
        paths: List[str] = []
        for j, (x, y) in enumerate(SMALL_POS):
            p = f"{env_root}/box_{j}"
            BOX_SMALL_CFG.func(
                p, BOX_SMALL_CFG,
                translation=(x, y, 1.0),
                orientation=self._quat_z(0.0, dev)
            )
            self._ensure_dynamic_rigid_body(p, mass=1.0)
            paths.append(p)
        self.small_box_paths[eid] = paths

    def _teleport_boxes(self, eid: int, dev: torch.device):
        big_p = (0.0, 0.75, 0.9)
        big_q = self._wxyz_to_xyzw(self._quat_z(math.pi / 2, dev))
        self._set_rb_state(self.big_box_path[eid], big_p, big_q)

        for j, (x, y) in enumerate(SMALL_POS):
            p_xyz = (x, y, 0.9)
            q_xyzw = self._wxyz_to_xyzw(self._quat_z(0.0, dev))
            self._set_rb_state(self.small_box_paths[eid][j], p_xyz, q_xyzw)

    # ------------------------------------------------------------------
    # Observation / Reward
    # ------------------------------------------------------------------
    def obs(self, eid: int) -> np.ndarray:
        return self.robot.data.joint_pos[eid].cpu().numpy().astype(np.float32)

    def reward(self, eid: int) -> float:
        big = self.big_xy[eid]
        small = self.small_xy[eid]
        return -float(np.mean(np.linalg.norm(big[None, :] - small, axis=1)))


# ───────────────────────── 6)  Env-Config & Env ───────────────────────────
@configclass
class ClearanceEnvCfg(DirectRLEnvCfg):
    num_envs: int = args_cli.num_envs
    decimation: int = 1
    episode_length_s: float = 30.0

    sim: sim_utils.SimulationCfg = sim_utils.SimulationCfg(
        use_fabric=False,
        gravity=(0.0, 0.0, -9.81)
    )

    scene: InteractiveSceneCfg = InteractiveSceneCfg(
        num_envs=num_envs,
        env_spacing=10.0,
        replicate_physics=True
        # filter_collisions=True ist Default; Boden ist pro Env, daher OK
    )

    usd_path, mount_path, robot_root = USD_ALL, MOUNT_PATH, ROBOT_ROOT
    robot_cfg: ArticulationCfg = ArticulationCfg(
        prim_path=robot_root,  # wird in _setup_scene auf ein Pattern ersetzt
        actuators={
            "all": ImplicitActuatorCfg(
                joint_names_expr=[".*"],
                stiffness=400.0,
                damping=40.0
            )
        },
    )

    observation_space = spaces.Box(-math.pi, math.pi, (DOF,), dtype=np.float32)
    action_space      = spaces.Box(-0.05,   0.05,  (DOF,), dtype=np.float32)
    DOF: int = DOF


class ClearanceEnv(DirectRLEnv):
    cfg: ClearanceEnvCfg

    # ---------- Szene aufsetzen -----------------------------------------
    def _setup_scene(self):
        # Licht
        sim_utils.DomeLightCfg(intensity=3000., color=(0.75,)*3).func(
            "/World/Light", sim_utils.DomeLightCfg()
        )

        # >>> Boden unter env_0 spawnen, damit er in alle Envs geklont wird <<<
        env0_root = "/World/envs/env_0"
        gp_cfg = sim_utils.GroundPlaneCfg()  # Ebene bei z=0 mit Collider
        gp_cfg.func(f"{env0_root}/ground", gp_cfg)

        # USD unter env_0 mounten, damit clone_environments ihn dupliziert
        spawn_from_usd(self.cfg.mount_path, UsdFileCfg(usd_path=self.cfg.usd_path))

        # ArticulationRoot an der env_0-Instanz sicherstellen
        stage = self.scene.stage
        prim  = stage.GetPrimAtPath(self.cfg.robot_root)
        if not UsdPhysics.ArticulationRootAPI(prim):
            UsdPhysics.ArticulationRootAPI.Apply(prim)

        # Articulation über ALLE Envs via Pattern-Pfad adressieren
        pattern_path = self.cfg.robot_root.replace("/env_0/", "/env_.*/")
        robot_cfg_all = self.cfg.robot_cfg.replace(prim_path=pattern_path)
        self.robot = Articulation(robot_cfg_all)

        # env_0 → env_1..N klonen (Physik replizieren!)
        self.scene.clone_environments(copy_from_source=True)
        self.sim.render()

        self.task = ClearanceTask(self.scene, self.robot, self.cfg.DOF)

        # ─── Episoden entkoppeln (Jitter) ─────────────────────────
        base_steps = int(self.cfg.episode_length_s / self.step_dt)
        self._horizon_jitter = 120
        self._ep_horizon = torch.full(
            (self.num_envs,),
            base_steps,
            dtype=torch.int32,
            device=self.device
        )

    # ---------- RL-Schleifen-Hooks --------------------------------------
    def _pre_physics_step(self, actions: torch.Tensor):
        self._actions = actions.clone()

    def _apply_action(self):
        # Batched DOF-Write
        dev = self.robot.data.joint_pos.device
        q_cur = self.robot.data.joint_pos            # (num_envs, DOF)
        q_tgt = torch.clamp(q_cur + self._actions.to(dev), -math.pi, math.pi)
        env_ids = torch.arange(self.num_envs, dtype=torch.int32, device=dev)
        vel = torch.zeros_like(q_tgt)
        self.robot.write_joint_state_to_sim(q_tgt, vel, env_ids=env_ids)

    def _get_observations(self):
        obs = self.robot.data.joint_pos
        return {"policy": obs}

    def _get_rewards(self):
        r = [self.task.reward(e) for e in range(self.num_envs)]
        return torch.tensor(r, device=self.device)

    def _get_dones(self):
        term = torch.zeros(self.num_envs, dtype=torch.bool, device=self.device)
        trunc = self.episode_length_buf >= (self._ep_horizon - 1)
        return term, trunc

    def _reset_idx(self, env_ids):
        dev = self.robot.data.joint_pos.device
        env_ids = env_ids.to(device=dev, dtype=torch.int32)

        self.task.reset_batch(env_ids)

        base_steps = int(self.cfg.episode_length_s / self.step_dt)
        j = self._horizon_jitter
        rand = torch.randint(
            low=-j, high=j + 1,
            size=(env_ids.shape[0],),
            device=self.device,
            dtype=torch.int32
        )
        self._ep_horizon[env_ids.long()] = torch.clamp(base_steps + rand, min=1)

# ───────────────────────── 7)  skrl-Modelle  ──────────────────────────────
class Policy(GaussianMixin, Model):
    """Stochastisches MLP-Policy-Net mit lernbarem log-σ"""
    def __init__(self, obs_space, act_space, device, **kw):
        Model.__init__(self, obs_space, act_space, device)
        GaussianMixin.__init__(self, **kw)
        self.net = nn.Sequential(
            nn.Linear(self.num_observations, 256), nn.Tanh(),
            nn.Linear(256, 256), nn.Tanh(),
            nn.Linear(256, self.num_actions)
        )
        self.log_std_parameter = nn.Parameter(torch.zeros(self.num_actions))

    def compute(self, inputs, role):
        return self.net(inputs["states"]), self.log_std_parameter, {}

class Value(DeterministicMixin, Model):
    """Deterministisches Wert-Net (MLP)"""
    def __init__(self, obs_space, act_space, device, **kw):
        Model.__init__(self, obs_space, act_space, device)
        DeterministicMixin.__init__(self, **kw)
        self.net = nn.Sequential(
            nn.Linear(self.num_observations, 256), nn.Tanh(),
            nn.Linear(256, 256), nn.Tanh(),
            nn.Linear(256, 1)
        )

    def compute(self, inputs, role):
        return self.net(inputs["states"]), {}

# ───────────────────────── 8)  Main  ──────────────────────────────────────
def main():
    print("[DEBUG] Starting environment configuration")
    env_cfg = ClearanceEnvCfg()
    print(f"[DEBUG] Created ClearanceEnvCfg: {env_cfg}")
    isaac_env = ClearanceEnv(env_cfg)
    print(f"[DEBUG] Instantiated ClearanceEnv: {isaac_env}")

    # Offizieller IsaacLab-Wrapper
    env = wrap_env(isaac_env, wrapper="isaaclab")
    print(f"[DEBUG] Wrapped env: {env} (obs_space={env.observation_space}, act_space={env.action_space})")
    print(f"[DEBUG] Observation space: {env.observation_space}, shape: {env.observation_space.shape}")

    # Seeds
    print("[DEBUG] Seeding environments and libraries")
    if SEED is not None:
        env.seed(SEED)
        torch.manual_seed(SEED)
        np.random.seed(SEED)
        random.seed(SEED)
        print(f"[DEBUG] Seeds set to {SEED}")

    # Modelle
    print("[DEBUG] Building policy and value networks")
    policy = Policy(env.observation_space, env.action_space,
                    device="cuda", clip_actions=True, clip_log_std=True)
    print(f"[DEBUG] Policy model: {policy}")
    value = Value(env.observation_space, env.action_space, device="cuda")
    print(f"[DEBUG] Value model: {value}")
    models = {"policy": policy, "value": value}

    # Speicher
    print("[DEBUG] Setting up memory buffer")
    memory = RandomMemory(
        memory_size=ROLLOUT_LENGTH, 
        num_envs=env.num_envs, 
        device=torch.device("cuda"))
    print(f"[DEBUG] Memory initialized: size={memory.memory_size}, num_envs={memory.num_envs}")

    # PPO-Agent
    print("[DEBUG] Configuring PPO agent")
    cfg = PPO_DEFAULT_CONFIG.copy()
    ckpt_every = ROLLOUT_LENGTH * env.num_envs  #327680
    cfg.update({
        "rollouts":           1024,
        "learning_epochs":    2,
        "mini_batches":       16,
        "mixed_precision":    True,
        "discount_factor":    0.99,
        "learning_rate":      LEARNING_RATE,
        "value_loss_scale":   0.5,
        "entropy_loss_scale": 0.01,
        "gradient_clipping":  0.5,
    })
    cfg["experiment"]["directory"] = str(pathlib.Path("checkpoints"))
    cfg["experiment"]["checkpoint_interval"] = ckpt_every
    print(f"[DEBUG] PPO config: {cfg}")

    agent = PPO(models=models, memory=memory, cfg=cfg, device="cuda")
    agent.env               = env
    agent.observation_space = env.observation_space
    agent.action_space      = env.action_space
    print("[DEBUG] PPO agent created.")

    # Trainer
    print("[DEBUG] Initializing trainer")
    trainer_cfg = {"timesteps": TOTAL_TIMESTEPS, "headless": args_cli.headless}
    print(f"[DEBUG] Trainer config: {trainer_cfg}, env.num_envs={env.num_envs}")
    trainer = SequentialTrainer(cfg=trainer_cfg, env=env, agents=agent)
    print("[DEBUG] Trainer instance created.")

    # Training
    print("[DEBUG] Starting training loop")
    trainer.train()
    print("[DEBUG] Training loop finished")

    agent.save("ppo_hc10_final")
    print("✅ Training complete – model saved as 'ppo_hc10_final'")

    # Cleanup
    print("[DEBUG] Closing environments and simulator")
    env.close()
    simulation_app.close()
    print("[DEBUG] Shutdown complete")

if __name__ == "__main__":
    main()

[RandomOakForest]

Thank you for starting this discussion. What is what you mean by physically correct? What is what you are observing?

[Icon45]

Hi,

with physically correct i mean that the boxes that i spawn should fall to the ground or onto the next rigid body. Currently they are just floating in mid air. (Sorry forgot to mention my observation in the original post)