Fetching the psoe of a rigid body seems bugged

[Icon45]

Hi erveryone, i am trying to track may spawned boxes. they are part of my env and can and should be manipulated by my robot. Based on how they are manipulated the robot should learn a specific behavior. however i am really not getting why i cannot track the boxes and there think root_state_w is bugged.

If its not a bug i would really apreaciat any help in getting it to work:

 -*- coding: utf-8 -*-
from __future__ import annotations

"""
Hc10Env – vollständige Version
- Tracking über root_state_w (statt USD)
- Reward:
  (1) Orangen weg von Blau (mean Distanz)
  (2) EE-Push: Bewegung (Geschwindigkeit) von Orangen in EE-Nähe (link_6_t)
  (3) Blau muss bleiben (starker Penalty bei Verschieben + Speed-Penalty)
  (4) Kein Bodenkontakt (Penalty über min Link-Z)
"""

import math
import numpy as np
import torch

from isaaclab.envs import DirectRLEnv
from isaaclab.assets import Articulation, RigidObject
from isaaclab.assets.rigid_object import RigidObjectCfg
import isaaclab.sim as sim_utils

import carb

# ----------------------------- Farben / Layout
BLUE, ORANGE = (0.2, 0.4, 1.0), (1.0, 0.5, 0.2)

SMALL_POS: tuple[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_Z_OFFSET = 0.05
BOX_Z_BASE   = 0.10
BOX_Z        = BOX_Z_BASE + BOX_Z_OFFSET

SMALL_BOX_WIDTH = 0.1

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=(SMALL_BOX_WIDTH, SMALL_BOX_WIDTH, SMALL_BOX_WIDTH),
    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),
)

# ----------------------------- Rewards / Dones
CLEARANCE_THRESHOLD = 0.30
MOVE_PRINT_THRESH   = 1e-4
VEL_PRINT_THRESH    = 1e-3

BLUE_BOX_QUAT = (0.70710678, 0.0, 0.0, 0.70710678)
IDENTITY_QUAT = (1.0, 0.0, 0.0, 0.0)

# === Reward / Safety Tuning ===
EE_LINK_NAME     = "link_6_t"   # Endeffektor-Link
EE_NEAR_RADIUS   = 0.15         # [m]
Z_GROUND_THRESH  = 0.02         # [m]

W_CLEAR          = 1.0          # Orangen weg von Blau
W_EE_PUSH        = 0.6          # Bewegung in EE-Nähe
L_BLUE_STAY      = 10.0          # Blau soll bleiben
L_NO_GROUND      = 5.0          # Kein Bodenkontakt


class Hc10Env(DirectRLEnv):
    """HC10-Task mit robustem Tracking und angepasstem Reward."""
    from .hc10_env_cfg import Hc10EnvCfg  # type: ignore[assignment]
    cfg: "Hc10EnvCfg"

    def __init__(self, cfg: "Hc10EnvCfg", render_mode: str | None = None, **kwargs):
        # Handles
        self.robot: Articulation | None = None
        self._dof: int | None = None
        self.box_big: RigidObject | None = None
        self.box_smalls: list[RigidObject] = []
        self._num_small = len(SMALL_POS)

        # Source env path
        self._src_env_path: str | None = None

        # Action/limits
        self.delta_clip = 0.05
        self.actions: torch.Tensor | None = None
        self._q_target: torch.Tensor | None = None
        self._lo: torch.Tensor | None = None
        self._hi: torch.Tensor | None = None

        # Tracking-Puffer
        self._usd_prev_big_p: np.ndarray | None = None     # [N, 3]
        self._usd_prev_small_p: np.ndarray | None = None   # [N, K, 3]
        self._usd_vel_big: np.ndarray | None = None        # [N, 3]
        self._usd_vel_small: np.ndarray | None = None      # [N, K, 3]

        # Episode-Tracker
        self._ep_len_steps: int | None = None
        self._progress_steps: torch.Tensor | None = None
        self._episode_returns: torch.Tensor | None = None
        self._episode_counts: torch.Tensor | None = None
        self._last_reward: torch.Tensor | None = None

        # Reward-Hilfen
        self._ee_body_idx: int | None = None
        self._blue_spawn_xy: torch.Tensor | None = None   # [N, 2]
        self._small_spawn_xy: torch.Tensor | None = None  # [N, K, 2]

        # Init-Maske
        self._env_inited: torch.Tensor | None = None

        super().__init__(cfg, render_mode, **kwargs)

    # -------------------------- Helpers --------------------------
    def _box_paths_for_env(self, eid: int) -> tuple[str, list[str]]:
        base = f"/World/envs/env_{eid}"
        big = f"{base}/box_big"
        smalls = [f"{base}/box_small_{k}" for k in range(self._num_small)]
        return big, smalls

    def _get_ee_index(self) -> int:
        if self._ee_body_idx is not None:
            return self._ee_body_idx
        names = list(self.robot.body_names)
        if EE_LINK_NAME in names:
            self._ee_body_idx = names.index(EE_LINK_NAME)
        else:
            carb.log_warn(f"[reward] EE link '{EE_LINK_NAME}' nicht gefunden – nutze letztes Body.")
            self._ee_body_idx = len(names) - 1
        return self._ee_body_idx

    # -------------------------- Tracking --------------------------
    def _update_tracking(self) -> None:
        """Tracking basierend auf root_state_w (Pos/Vel), + Logging + extras."""
        N = self.num_envs
        K = self._num_small
        dt = float(self.cfg.sim.dt) * float(self.cfg.decimation)
        dt_clamped = max(dt, 1e-8)

        # Big box
        big_state   = self.box_big.data.root_state_w       # [N, 13]
        big_now_pos = big_state[:, 0:3].cpu().numpy()
        big_now_vel = big_state[:, 7:10].cpu().numpy()

        # Small boxes
        small_now_pos = np.zeros((N, K, 3), dtype=np.float64)
        small_now_vel = np.zeros((N, K, 3), dtype=np.float64)
        for i, box in enumerate(self.box_smalls):
            st = box.data.root_state_w
            small_now_pos[:, i] = st[:, 0:3].cpu().numpy()
            small_now_vel[:, i] = st[:, 7:10].cpu().numpy()

        # Displacements
        d_big   = big_now_pos - self._usd_prev_big_p
        d_small = small_now_pos - self._usd_prev_small_p
        v_big_est   = d_big / dt_clamped
        v_small_est = d_small / dt_clamped

        n_big    = np.linalg.norm(d_big, axis=1)
        n_small  = np.linalg.norm(d_small, axis=2)
        vn_big   = np.linalg.norm(big_now_vel, axis=1)
        vn_small = np.linalg.norm(small_now_vel, axis=2)

        # Logging bei signifikanter Bewegung
        for e in range(N):
            if np.isfinite(n_big[e]) and (n_big[e] > MOVE_PRINT_THRESH or vn_big[e] > VEL_PRINT_THRESH):
                px, py, pz = big_now_pos[e]
                carb.log_warn(
                    f"box_big env {e}: Δ={n_big[e]:.5e} m, |v|={vn_big[e]:.5e} m/s @ ({px:.3f}, {py:.3f}, {pz:.3f})"
                )
            for j in range(K):
                if np.isfinite(n_small[e, j]) and (n_small[e, j] > MOVE_PRINT_THRESH or vn_small[e, j] > VEL_PRINT_THRESH):
                    px, py, pz = small_now_pos[e, j]
                    carb.log_warn(
                        f"box_small_{j} env {e}: Δ={n_small[e, j]:.5e} m, |v|={vn_small[e, j]:.5e} m/s @ ({px:.3f}, {py:.3f}, {pz:.3f})"
                    )

        # extras
        self.extras["track/box_big_pos_w"]   = torch.from_numpy(big_now_pos).float().cpu()
        self.extras["track/box_small_pos_w"] = torch.from_numpy(small_now_pos).float().cpu()
        self.extras["track/box_big_vel_w"]   = torch.from_numpy(big_now_vel).float().cpu()
        self.extras["track/box_small_vel_w"] = torch.from_numpy(small_now_vel).float().cpu()
        self.extras["track/box_big_dnorm"]   = torch.from_numpy(n_big).float().view(-1, 1).cpu()
        self.extras["track/box_small_dnorm"] = torch.from_numpy(n_small).float().cpu()
        self.extras["track/box_small_dnorm_mean"] = torch.from_numpy(np.mean(n_small, axis=1, keepdims=True)).float().cpu()
        self.extras["track/box_small_dnorm_max"]  = torch.from_numpy(np.max(n_small, axis=1, keepdims=True)).float().cpu()
        self.extras["track/box_big_vnorm"]        = torch.from_numpy(vn_big).float().view(-1, 1).cpu()
        self.extras["track/box_small_vnorm_mean"] = torch.from_numpy(np.mean(vn_small, axis=1, keepdims=True)).float().cpu()
        self.extras["track/box_small_vnorm_max"]  = torch.from_numpy(np.max(vn_small, axis=1, keepdims=True)).float().cpu()

        # Update prev
        self._usd_prev_big_p   = big_now_pos
        self._usd_prev_small_p = small_now_pos
        self._usd_vel_big      = big_now_vel
        self._usd_vel_small    = small_now_vel

    # -------------------------- Scene --------------------------
    def _setup_scene(self):
        self.robot = self.scene["robot"]
        self._src_env_path = "/World/envs/env_0"

        # Spawn große Box
        self.box_big = sim_utils.spawn_cuboid(
            cfg=BOX_BIG_CFG,
            prim_path=f"{self._src_env_path}/box_big",
            translation=(0.0, 0.75, BOX_Z),
            orientation=BLUE_BOX_QUAT,
        )
        # Spawn kleine Boxen
        for i, (x, y) in enumerate(SMALL_POS):
            sim_utils.spawn_cuboid(
                cfg=BOX_SMALL_CFG,
                prim_path=f"{self._src_env_path}/box_small_{i}",
                translation=(x, y, BOX_Z),
                orientation=IDENTITY_QUAT,
            )

        # Envs klonen
        self.scene.clone_environments(copy_from_source=True)

        # RigidObject-Views
        self.box_big = RigidObject(cfg=RigidObjectCfg(prim_path="/World/envs/env_.*/box_big"))
        self.box_smalls = [
             # Spawn kleine Boxen
            for i, (x, y) in enumerate(SMALL_POS):
                sim_utils.spawn_cuboid(
                    cfg=BOX_SMALL_CFG,
                    prim_path=f"{self._src_env_path}/box_small_{i}",
                    translation=(x, y, BOX_Z),
                    orientation=IDENTITY_QUAT,
                )
        ]

    # -------------------------- Reset/Obs --------------------------
    def _reset_idx(self, env_ids: torch.Tensor):
        if self.robot is None:
            self.robot = self.scene.articulations.get("robot", next(iter(self.scene.articulations.values())))

        env_ids = env_ids.to(dtype=torch.int64, device=self.device)

        # DoF/Limits
        if self._dof is None:
            self._dof = int(self.robot.data.joint_pos.shape[1])
            self._q_target = torch.zeros((self.num_envs, self._dof), dtype=torch.float32, device=self.device)
            limits = getattr(self.cfg.scene.robot, "joint_limits", None)
            if not limits:
                lo = [-math.pi] * self._dof
                hi = [ math.pi] * self._dof
            else:
                names = list(self.robot.dof_names)
                lo = [limits.get(n, (-math.pi, math.pi))[0] for n in names]
                hi = [limits.get(n, (-math.pi, math.pi))[1] for n in names]
            self._lo = torch.tensor(lo, dtype=torch.float32, device=self.device).expand(self.num_envs, -1)
            self._hi = torch.tensor(hi, dtype=torch.float32, device=self.device).expand(self.num_envs, -1)

        # Episodenlänge (Steps)
        self._ep_len_steps = max(1, int(math.ceil(float(self.cfg.episode_length_s) /
                              (float(self.cfg.sim.dt) * float(self.cfg.decimation)))))

        # Puffer allozieren
        if self._progress_steps is None:
            self._progress_steps  = torch.zeros(self.num_envs, dtype=torch.int32, device=self.device)
            self._episode_returns = torch.zeros(self.num_envs, dtype=torch.float32, device=self.device)
            self._episode_counts  = torch.zeros(self.num_envs, dtype=torch.int32, device=self.device)
            self._last_reward     = torch.zeros(self.num_envs, dtype=torch.float32, device=self.device)

        if self._env_inited is None:
            self._env_inited = torch.zeros(self.num_envs, dtype=torch.bool, device=self.device)

        # --- Roboter zurücksetzen ---
        root = self.robot.data.default_root_state[env_ids].clone()
        root[:, :3] = self.scene.env_origins[env_ids]
        self.robot.write_root_pose_to_sim(root[:, :7])
        self.robot.write_root_velocity_to_sim(root[:, 7:])
        qz = torch.zeros((len(env_ids), self._dof), dtype=torch.float32, device=self.device)
        self.robot.write_joint_state_to_sim(qz, torch.zeros_like(qz))
        self._q_target[env_ids] = 0.0

        # --- Boxen zurücksetzen (PhysX) ---
        blue_pose = torch.zeros((len(env_ids), 7), dtype=torch.float32, device=self.device)
        blue_vel  = torch.zeros((len(env_ids), 6), dtype=torch.float32, device=self.device)
        blue_pose[:, 0:3] = self.scene.env_origins[env_ids] + torch.tensor((0.0, 0.75, BOX_Z), device=self.device)
        blue_pose[:, 3:7] = torch.tensor(BLUE_BOX_QUAT, dtype=torch.float32, device=self.device)
        self.box_big.write_root_pose_to_sim(blue_pose)
        self.box_big.write_root_velocity_to_sim(blue_vel)

        for i, (x, y) in enumerate(SMALL_POS):
            pose = torch.zeros((len(env_ids), 7), dtype=torch.float32, device=self.device)
            vel  = torch.zeros((len(env_ids), 6), dtype=torch.float32, device=self.device)
            pose[:, 0:3] = self.scene.env_origins[env_ids] + torch.tensor((x, y, BOX_Z), device=self.device)
            pose[:, 3:7] = torch.tensor(IDENTITY_QUAT, dtype=torch.float32, device=self.device)
            self.box_smalls[i].write_root_pose_to_sim(pose)
            self.box_smalls[i].write_root_velocity_to_sim(vel)

        self.scene.write_data_to_sim()

        # === Tracking-Reset ===
        try:
            self.scene.read_data_from_sim()
        except Exception:
            pass

        N, K = self.num_envs, self._num_small
        big_state = self.box_big.data.root_state_w
        big_now = big_state[:, 0:3].cpu().numpy()
        small_now = np.zeros((N, K, 3), dtype=np.float64)
        for i, box in enumerate(self.box_smalls):
            state = box.data.root_state_w
            small_now[:, i] = state[:, 0:3].cpu().numpy()
        self._usd_prev_big_p   = big_now
        self._usd_prev_small_p = small_now
        self._usd_vel_big      = np.zeros_like(big_now)
        self._usd_vel_small    = np.zeros_like(small_now)

        # Spawn-XYs für Reward cachen
        self._blue_spawn_xy = self.box_big.data.root_state_w[:, 0:2].detach().clone()
        small_xy_list = [b.data.root_state_w[:, 0:2] for b in self.box_smalls]
        self._small_spawn_xy = torch.stack(small_xy_list, dim=1).detach().clone()  # [N, K, 2]

        self._progress_steps[env_ids]  = 0
        self._episode_returns[env_ids] = 0.0

        self.scene.reset(env_ids=env_ids)
        self._env_inited[env_ids] = True

    def _get_observations(self) -> dict:
        q = self.robot.data.joint_pos  # [N, DoF]
        N, K = self.num_envs, self._num_small
        xy_feat = torch.zeros((N, 2 * (1 + K)), dtype=torch.float32, device=self.device)

        big_state = self.box_big.data.root_state_w
        b_xy = big_state[:, 0:2]
        xy_feat[:, 0:2] = b_xy
        off = 2
        for j, box in enumerate(self.box_smalls):
            s_xy = box.data.root_state_w[:, 0:2]
            xy_feat[:, off + 2*j : off + 2*(j+1)] = s_xy
        return {"policy": torch.cat([q, xy_feat], dim=-1)}

    def _pre_physics_step(self, actions: torch.Tensor | None):
        try:
            self.scene.read_data_from_sim()
        except Exception:
            pass
        self._update_tracking()

        if actions is None:
            if self.actions is None:
                self.actions = torch.zeros_like(self._q_target, device=self.device)
            return
        a = torch.as_tensor(actions, device=self.device, dtype=torch.float32).reshape(self._q_target.shape)
        self.actions = torch.clamp(a, -self.delta_clip, self.delta_clip)

    def _apply_action(self):
        q = self.robot.data.joint_pos
        q_new = torch.clamp(q + self.actions, min=self._lo, max=self._hi)
        self._q_target = q_new
        self.robot.set_joint_position_target(self._q_target)
        self.scene.write_data_to_sim()

    # -------------------------- Rewards / Dones --------------------------
    def _get_rewards(self):
        try:
            self.scene.read_data_from_sim()
        except Exception:
            pass
        self._update_tracking()

        N, K = self.num_envs, self._num_small
        device = self.device

        # Big box
        big_state   = self.box_big.data.root_state_w               # [N, 13]
        big_pos_xy  = big_state[:, 0:2]                            # [N, 2]
        big_speed   = torch.linalg.norm(big_state[:, 7:10], dim=-1)

        # Small boxes
        small_pos   = torch.stack([b.data.root_state_w[:, 0:3] for b in self.box_smalls], dim=1)  # [N, K, 3]
        small_xy    = small_pos[..., :2]                                                               # [N, K, 2]
        small_vel   = torch.stack([b.data.root_state_w[:, 7:10] for b in self.box_smalls], dim=1)  # [N, K, 3]
        small_speed = torch.linalg.norm(small_vel, dim=-1)                                         # [N, K]

        # (1) Orangen weg von Blau
        dist_small_to_blue = torch.linalg.norm(small_xy - big_pos_xy[:, None, :], dim=-1)  # [N, K]
        r_clear = dist_small_to_blue.mean(dim=1)                                           # [N]

        # (2) EE-Push (Bewegung in EE-Nähe)
        ee_idx  = self._get_ee_index()
        body_st = self.robot.data.body_state_w                    # [N, B, 13]
        ee_pos  = body_st[:, ee_idx, 0:3]                         # [N, 3]
        dist_ee = torch.linalg.norm(small_pos - ee_pos[:, None, :], dim=-1)  # [N, K]
        near_ee = (dist_ee < EE_NEAR_RADIUS).float()
        r_ee_push = (near_ee * small_speed).mean(dim=1)           # [N]

        # (3) Blau soll bleiben
        if self._blue_spawn_xy is None:
            self._blue_spawn_xy = big_pos_xy.detach().clone()
        blue_disp = torch.linalg.norm(big_pos_xy - self._blue_spawn_xy, dim=-1)
        r_blue_stay = -L_BLUE_STAY * blue_disp - 0.25 * L_BLUE_STAY * big_speed

        # (4) Kein Bodenkontakt
        all_body_pos = body_st[:, :, 0:3]      # [N, B, 3]
        min_z, _     = all_body_pos[..., 2].min(dim=1)
        r_no_ground  = -L_NO_GROUND * torch.relu(Z_GROUND_THRESH - min_z) / Z_GROUND_THRESH

        # Gesamt
        reward = W_CLEAR * r_clear + W_EE_PUSH * r_ee_push + r_blue_stay + r_no_ground

        # Extras
        self.extras["reward/clear_mean_dist"]   = r_clear.detach().cpu()
        self.extras["reward/ee_push"]           = r_ee_push.detach().cpu()
        self.extras["reward/blue_disp_xy"]      = blue_disp.detach().cpu()
        self.extras["reward/min_link_z"]        = min_z.detach().cpu()
        self.extras["reward/total"]             = reward.detach().cpu()

        self._last_reward = reward
        if self._episode_returns is None:
            self._episode_returns = torch.zeros_like(reward)
        self._episode_returns += reward
        return reward

    def _is_done(self):
        self.done = torch.zeros(self.num_envs, dtype=torch.bool, device=self.device)

    def _get_dones(self):
        terms = []
        big_state = self.box_big.data.root_state_w
        big_xy = big_state[:, 0:2].cpu().numpy()
        for e in range(self.num_envs):
            b = big_xy[e]
            dmin = float("inf")
            for box in self.box_smalls:
                s = box.data.root_state_w[e, 0:2].cpu().numpy()
                d = float(np.linalg.norm(s - b))
                if d < dmin:
                    dmin = d
            terms.append(dmin >= CLEARANCE_THRESHOLD)
        terminated = torch.tensor(terms, dtype=torch.bool, device=self.device)

        if self._progress_steps is None:
            self._progress_steps = torch.zeros(self.num_envs, dtype=torch.int32, device=self.device)
        self._progress_steps += 1
        trunc = self._progress_steps >= int(self._ep_len_steps or 1)

        if self._episode_counts is None:
            self._episode_counts = torch.zeros(self.num_envs, dtype=torch.int32, device=self.device)
        ended = terminated | trunc
        self._episode_counts[ended] += 1

        self.extras["episode_step"] = self._progress_steps.clone().cpu()
        self.extras["episode_index"] = self._episode_counts.clone().cpu()
        self.extras["diag/clearance_threshold"] = torch.tensor([CLEARANCE_THRESHOLD], dtype=torch.float32)
        return terminated, trunc

[RandomOakForest]

Thank you for posting this. If you could share an error log and your tool versions would help reproduce your issue. For now, I'll move your post to our Discussions for follow up. Here are a few things to consider.

The root cause of difficulties tracking and manipulating spawned rigid boxes in Isaac Lab commonly relates to the update timing and management of simulation buffers, especially for the root_state_w after resets or state changes. Isaac Lab uses the root state buffer in the simulation world frame, and this buffer is updated only when simulation steps are correctly executed or explicit update methods are called.

Correct approach to tracking boxes

The positions and orientations of each spawned box (rigid object) must be accessed and updated through the simulation buffers, not just the asset data. The recommended workflow is:

• After spawning or moving a rigid object, call write_root_state_to_sim() on the object, passing the new position and velocity in world coordinates.
• Immediately call the object’s reset() method to ensure internal buffers reflect the changes before stepping the simulation.
• Step the simulation using the environment’s step call. Only after this step will the updated states (including root_state_w) be accurately reflected in buffers used for observation and learning.

Common pitfalls and Isaac Lab-specific bugs

A known issue with Isaac Lab (as documented in its official troubleshooting and known issues pages) is the presence of stale values in root_state_w and related buffers, especially after resets. This occurs because PhysX (the underlying physics engine) updates certain states only after a full simulation step; writing new states to the simulation buffer does not immediately update all observables until the simulation progresses.

• Simply resetting or setting the pose/velocity is not sufficient; the buffers will contain outdated data until the next simulation tick.
• If manipulating objects for RL or learning, always ensure your policy sampling or observations occur after stepping the simulation.

Steps for robust box tracking

1. After manipulating a box (directly or via robot), call write_root_state_to_sim(new_state) on all relevant box objects.
2. Call reset() on each rigid object to clear internal simulation data.
3. Step the simulation environment (env.step() or similar).
4. Now query root_state_w for each box for the current, correct position/rotation data.

This workflow avoids the stale buffer issue and lets the RL agent receive fresh observations after every environment interaction. A feature request exists for finer-grained control over PhysX state and buffer updates, but as of Isaac Lab 1.2+, this workaround remains necessary.