marl-world-model / env_multi_agent_isaac.py
Devin
MARL + world model source code and design docs (2026-08-13 03:14:28)
33ffc04 3 months ago
43.6 kB
r"""
Multi-Agent Isaac Sim environment (headless, D3D12/Vulkan-workaround) for
Isaac Sim 5.1.0.
Run the smoke test with:
    H:\Robotics\ISSAC_SIM_5.1.0\python.bat H:\Robotics\marl_world_model\env_multi_agent_isaac.py --smoke
The environment spawns a small warehouse floor, two mobile robots (differential
or primitive box-based fallback), colour-coded cubes and bins.  It exposes the
standard RL interface:
    __init__, reset(), step(actions), get_observations(), get_rewards(), render()
Observations:  RGB 128x128 camera image + low-dim state
                (x, y, yaw, vx, vy, omega, nearest-cube dx/dy/dist, target-bin dx/dy/dist).
Actions:        continuous (num_agents, 4) = [forward, lateral, angular, gripper] in [-1, 1].
                gripper > 0.5 activates a simplified sticky gripper that grabs the
                nearest cube when the robot is close enough.
Reward:         +10 for a cube in the correct-colour bin
                -5  for a cube in the wrong-colour bin
                -1  robot-robot collision / out-of-bounds
                -0.1 step penalty
                + potential-based shaping for robot→cube and cube→bin approach.
If the env var SAVE_STAGE is set, the USD stage is saved to
H:\Robotics\marl_world_model\outputs\ on every reset.
"""

import os
import sys
import math
import time
import traceback

import numpy as np

# -----------------------------------------------------------------------------
# 1. Strip script-only arguments before launching the Isaac Sim app so the kit
#    launcher does not see unknown flags.
# -----------------------------------------------------------------------------
SMOKE_FLAG = "--smoke"
_IS_SMOKE = SMOKE_FLAG in sys.argv
if _IS_SMOKE:
    sys.argv.remove(SMOKE_FLAG)

# -----------------------------------------------------------------------------
# 2. Launch Isaac Sim FIRST.  All Omniverse-level imports must come after this.
# -----------------------------------------------------------------------------
from isaacsim import SimulationApp

SIMULATION_APP = SimulationApp(
    launch_config={
        "headless": True,
        "multi_gpu": False,
        "extra_args": [
            "--/app/vulkan=false",
            "--/renderer/multiGpu/enabled=false",
        ],
    },
    experience="",  # auto-detect the experience file
)

# -----------------------------------------------------------------------------
# 3. Omniverse / Isaac Sim imports (after SimulationApp)
# -----------------------------------------------------------------------------
from isaacsim.core.api import World
from isaacsim.core.api.objects import DynamicCuboid, FixedCuboid
from isaacsim.core.utils.numpy.rotations import euler_angles_to_quats
from isaacsim.sensors.camera import Camera
from isaacsim.storage.native import get_assets_root_path

try:
    from isaacsim.core.utils.prims import delete_prim, is_prim_path_valid

    _HAS_PRIM_UTILS = True
except Exception:
    _HAS_PRIM_UTILS = False

try:
    from isaacsim.robot.wheeled_robots.robots import WheeledRobot
    from isaacsim.robot.wheeled_robots.controllers.differential_controller import (
        DifferentialController,
    )
    from isaacsim.robot.wheeled_robots.controllers.holonomic_controller import (
        HolonomicController,
    )

    _HAS_WHEELED_ROBOTS = True
except Exception as _import_err:
    print("[WARN] Could not import WheeledRobot controllers:", _import_err)
    _HAS_WHEELED_ROBOTS = False

# -----------------------------------------------------------------------------
# Constants
# -----------------------------------------------------------------------------
NUM_AGENTS = 2
CAM_RES = (128, 128)
BASE_DIR = os.path.dirname(os.path.abspath(__file__))
OUTPUT_DIR = os.path.join(BASE_DIR, "outputs")


# -----------------------------------------------------------------------------
# Helper functions
# -----------------------------------------------------------------------------
def _yaw_from_quat(quat):
    """Scalar-first quaternion [w, x, y, z] -> yaw (radians)."""
    w, x, y, z = quat
    return math.atan2(2.0 * (w * z + x * y), 1.0 - 2.0 * (y * y + z * z))


class MultiAgentIsaacEnv:
    """Multi-agent warehouse-cube-sorting environment for Isaac Sim 5.1.0."""

    def __init__(
        self,
        num_agents=NUM_AGENTS,
        try_assets=True,
        max_linear_speed=0.3,
        max_angular_speed=1.0,
        floor_boundary=2.4,
    ):
        """
        Args:
            num_agents: Number of robots (default 2).
            try_assets: If True, attempt to load Jetbot from local Isaac Sim assets;
                        falls back to a primitive box robot if not available.
            max_linear_speed:  Action scaling in m/s.
            max_angular_speed: Action scaling in rad/s.
            floor_boundary:    Boundary used for out-of-bounds penalty.
        """
        self.num_agents = num_agents
        self.try_assets = try_assets
        self.max_linear_speed = float(max_linear_speed)
        self.max_angular_speed = float(max_angular_speed)
        self.floor_boundary = float(floor_boundary)
        self._dt = 1.0 / 60.0
        self._step_count = 0
        self.action_dim = 4  # [forward, lateral, angular, gripper]

        # Reward-shaping constants
        self._gamma = 0.99
        self._shaping_scale = 5.0       # cube→bin approach scaling
        self._approach_scale = 2.0      # robot→cube approach scaling
        self._gripper_threshold = 0.5   # gripper action > this activates grab
        self._grab_distance = 0.35      # max robot-cube distance for grabbing (m)
        self._bin_threshold = 0.25      # cube-bin distance for "in bin" (m)

        # ----------------------------------------------------------------------
        # Create the simulation world
        # ----------------------------------------------------------------------
        self._world = World(stage_units_in_meters=1.0)
        try:
            self._world.set_simulation_dt(
                physics_dt=self._dt, rendering_dt=self._dt
            )
        except Exception:
            # Older builds may not expose set_simulation_dt
            pass

        self._world.scene.add_default_ground_plane()

        # ----------------------------------------------------------------------
        # Scene: warehouse floor and lights
        # ----------------------------------------------------------------------
        self._spawn_floor_and_lights()

        # ----------------------------------------------------------------------
        # Scene: bins and cubes
        # ----------------------------------------------------------------------
        self._spawn_bins_and_cubes()

        # ----------------------------------------------------------------------
        # Scene: robots
        # ----------------------------------------------------------------------
        self._robot_type = self._spawn_robots(num_agents, try_assets)

        # ----------------------------------------------------------------------
        # Onboard RGB cameras for each agent
        # ----------------------------------------------------------------------
        self._cameras = self._spawn_cameras()

        # ----------------------------------------------------------------------
        # Reset to initialise physics and produce an initial observation
        # ----------------------------------------------------------------------
        obs = self.reset()

        # Expose spaces in a trainer-friendly format (per-agent list of dicts)
        self.observation_spaces = [
            {
                "image": {
                    "shape": tuple(obs["rgb"].shape[1:]),
                    "dtype": str(obs["rgb"].dtype),
                },
                "state": {
                    "shape": tuple(obs["state"].shape[1:]),
                    "dtype": str(obs["state"].dtype),
                },
            }
            for _ in range(self.num_agents)
        ]
        self.action_spaces = [
            {"shape": (4,), "low": -1.0, "high": 1.0}
            for _ in range(self.num_agents)
        ]

    # --------------------------------------------------------------------------
    # Scene construction helpers
    # --------------------------------------------------------------------------
    def _spawn_bins_and_cubes(self):
        """Create the three coloured bins and three coloured cubes."""
        self._bin_info = []
        self._cubes = []
        self._cube_info = []

        # Bin definitions: name, colour (RGB 0-1), and position
        bin_defs = [
            ("red", np.array([1.0, 0.0, 0.0]), np.array([2.0, -1.0, 0.02])),
            ("green", np.array([0.0, 1.0, 0.0]), np.array([2.0, 0.0, 0.02])),
            ("blue", np.array([0.0, 0.0, 1.0]), np.array([2.0, 1.0, 0.02])),
        ]

        for i, (name, colour, pos) in enumerate(bin_defs):
            bin_cuboid = self._world.scene.add(
                FixedCuboid(
                    prim_path=f"/World/bin_{name}",
                    name=f"bin_{name}",
                    position=pos,
                    size=0.4,
                    scale=np.array([1.0, 1.0, 0.1]),
                    color=colour,
                )
            )
            self._bin_info.append(
                {
                    "name": name,
                    "color_name": name,
                    "position": pos,
                    "prim": bin_cuboid,
                    "colour": colour,
                }
            )

        # Cube definitions
        cube_defs = [
            ("red", np.array([-0.5, 0.5, 0.05])),
            ("green", np.array([0.0, 0.5, 0.05])),
            ("blue", np.array([0.5, 0.5, 0.05])),
        ]
        colour_by_name = {n: c for n, c, _ in bin_defs}

        for i, (name, pos) in enumerate(cube_defs):
            colour = colour_by_name[name]
            cube = self._world.scene.add(
                DynamicCuboid(
                    prim_path=f"/World/cube_{i}",
                    name=f"cube_{i}",
                    position=pos,
                    size=0.08,
                    scale=np.array([1.0, 1.0, 1.0]),
                    color=colour,
                    mass=0.1,
                )
            )
            self._cubes.append(cube)
            self._cube_info.append(
                {"name": name, "position": pos, "prim": cube}
            )

    def _spawn_floor_and_lights(self):
        """Add a visible warehouse floor and lights so the cameras render colour."""
        try:
            self._world.scene.add(
                FixedCuboid(
                    prim_path="/World/warehouse_floor",
                    name="warehouse_floor",
                    position=np.array([0.0, 0.0, 0.0]),
                    size=10.0,
                    scale=np.array([1.0, 1.0, 0.02]),
                    color=np.array([0.3, 0.3, 0.32]),
                )
            )
        except Exception as e:
            print("[WARN] Could not add warehouse floor:", e)

        try:
            import omni.usd
            from pxr import UsdLux, Gf

            stage = omni.usd.get_context().get_stage()
            if stage:
                # Dome light for ambient illumination
                dome = UsdLux.DomeLight.Define(stage, "/World/DomeLight")
                dome.CreateColorAttr().Set(Gf.Vec3f(1.0, 1.0, 1.0))
                dome.CreateIntensityAttr().Set(800.0)
                dome.CreateExposureAttr().Set(0.0)

                # Directional light to cast shading and make objects visible
                distant = UsdLux.DistantLight.Define(stage, "/World/DistantLight")
                distant.CreateColorAttr().Set(Gf.Vec3f(1.0, 1.0, 1.0))
                distant.CreateIntensityAttr().Set(1500.0)
                distant.CreateExposureAttr().Set(0.0)
                distant.CreateAngleAttr().Set(1.0)
                # Point the light diagonally across the warehouse
                distant.AddRotateXOp().Set(45.0)
                distant.AddRotateYOp().Set(-30.0)
        except Exception as e:
            print("[WARN] Could not add lights:", e)

    def _spawn_robots(self, num_agents, try_assets):
        """Create the requested number of robots, falling back to box robots."""
        self._robots = []
        self._controllers = []

        # Default spawn positions on the left side of the warehouse
        robot_positions = [
            np.array([-1.5, 0.0, 0.05]),
            np.array([-1.5, -1.0, 0.05]),
            np.array([-1.5, 1.0, 0.05]),
            np.array([-1.5, -1.75, 0.05]),
        ][:num_agents]

        # ----------------------------------------------------------------------
        # Try to find local Isaac Sim assets
        # ----------------------------------------------------------------------
        assets_root = None
        if try_assets and _HAS_WHEELED_ROBOTS:
            try:
                assets_root = get_assets_root_path(skip_check=True)
                if assets_root:
                    print("[INFO] Isaac Sim assets root:", assets_root)
            except Exception as e:
                print("[INFO] get_assets_root_path returned:", e)
                assets_root = None

        # Treat the root as local only if it looks like a filesystem path
        is_local = bool(assets_root) and "://" not in assets_root

        # ----------------------------------------------------------------------
        # First attempt: holonomic Kaya (best matches the [fwd, lat, yaw] action)
        # ----------------------------------------------------------------------
        if is_local:
            kaya_path = os.path.join(
                assets_root, "Isaac", "Robots", "NVIDIA", "Kaya", "kaya.usd"
            )
            kaya_path = kaya_path.replace("\\", "/")
            if os.path.exists(kaya_path):
                try:
                    for i in range(num_agents):
                        robot = self._world.scene.add(
                            WheeledRobot(
                                prim_path=f"/World/agent_{i}",
                                name=f"agent_{i}",
                                wheel_dof_names=[
                                    "axle_0_joint",
                                    "axle_1_joint",
                                    "axle_2_joint",
                                ],
                                create_robot=True,
                                usd_path=kaya_path,
                                position=robot_positions[i],
                            )
                        )
                        self._robots.append(robot)

                    holonomic = HolonomicController(
                        "kaya_holonomic",
                        wheel_radius=np.array([0.04, 0.04, 0.04], dtype=np.float32),
                        wheel_positions=np.array(
                            [
                                [-0.0980432, 0.000636773, -0.050501],
                                [0.0493475, -0.084525, -0.050501],
                                [0.0495291, 0.0856937, -0.050501],
                            ],
                            dtype=np.float32,
                        ),
                        wheel_orientations=np.array(
                            [
                                [0.0, 0.0, 0.0, 1.0],
                                [0.866, 0.0, 0.0, -0.5],
                                [0.866, 0.0, 0.0, 0.5],
                            ],
                            dtype=np.float32,
                        ),
                        mecanum_angles=np.array([90, 90, 90], dtype=np.float32),
                        max_linear_speed=self.max_linear_speed,
                        max_angular_speed=self.max_angular_speed,
                    )
                    self._controllers = [holonomic] * num_agents
                    print("[INFO] Loaded holonomic Kaya robots.")
                    return "kaya"
                except Exception as e:
                    print("[WARN] Kaya load failed, will try Jetbot:", e)
                    self._clear_partial_robots(num_agents)
                    self._robots = []

            # ------------------------------------------------------------------
            # Second attempt: differential Jetbot (lateral action is ignored)
            # ------------------------------------------------------------------
            jetbot_path = os.path.join(
                assets_root, "Isaac", "Robots", "NVIDIA", "Jetbot", "jetbot.usd"
            )
            jetbot_path = jetbot_path.replace("\\", "/")
            if os.path.exists(jetbot_path):
                try:
                    for i in range(num_agents):
                        robot = self._world.scene.add(
                            WheeledRobot(
                                prim_path=f"/World/agent_{i}",
                                name=f"agent_{i}",
                                wheel_dof_names=[
                                    "left_wheel_joint",
                                    "right_wheel_joint",
                                ],
                                create_robot=True,
                                usd_path=jetbot_path,
                                position=robot_positions[i],
                            )
                        )
                        self._robots.append(robot)
                        self._controllers.append(
                            DifferentialController(
                                f"diff_ctrl_{i}",
                                wheel_radius=0.03,
                                wheel_base=0.1125,
                                max_linear_speed=self.max_linear_speed,
                                max_angular_speed=self.max_angular_speed,
                            )
                        )
                    print("[INFO] Loaded differential Jetbot robots.")
                    return "jetbot"
                except Exception as e:
                    print("[WARN] Jetbot load failed, using primitive robots:", e)
                    self._clear_partial_robots(num_agents)
                    self._robots = []

        # ----------------------------------------------------------------------
        # Fallback: driveable primitive boxes (3-DOF, fully supports [fwd, lat, yaw])
        # ----------------------------------------------------------------------
        print("[INFO] Using primitive box robots.")
        for i in range(num_agents):
            prim_path = f"/World/agent_{i}"
            if _HAS_PRIM_UTILS and is_prim_path_valid(prim_path):
                # In case an earlier failed attempt left a prim behind
                delete_prim(prim_path)

            robot = self._world.scene.add(
                DynamicCuboid(
                    prim_path=f"/World/agent_{i}_box",
                    name=f"agent_{i}",
                    position=robot_positions[i],
                    size=0.2,
                    scale=np.array([1.0, 0.8, 0.4]),
                    color=np.array([0.7, 0.7, 0.7]),
                    mass=2.0,
                )
            )
            self._robots.append(robot)

        self._controllers = [None] * num_agents
        return "primitive"

    def _clear_partial_robots(self, num_agents):
        """Remove any leftover prims from a failed WheeledRobot attempt."""
        if not _HAS_PRIM_UTILS:
            return
        for i in range(num_agents):
            for path in (f"/World/agent_{i}", f"/World/agent_{i}_box"):
                if is_prim_path_valid(path):
                    try:
                        delete_prim(path)
                    except Exception:
                        pass

    def _spawn_cameras(self):
        """Create an onboard RGB camera for each agent."""
        cameras = []
        for i, robot in enumerate(self._robots):
            robot_path = getattr(robot, "prim_path", f"/World/agent_{i}")
            cam = Camera(
                prim_path=f"{robot_path}/camera",
                name=f"agent_{i}_cam",
                frequency=60,
                resolution=CAM_RES,
                translation=np.array([0.0, 0.0, 0.6]),
                orientation=euler_angles_to_quats(
                    np.array([-15.0, -90.0, 0.0]), degrees=True
                ),
            )
            cameras.append(cam)
        return cameras

    # --------------------------------------------------------------------------
    # Public API
    # --------------------------------------------------------------------------
    def reset(self, seed=None):
        """Reset the environment and return the initial observation."""
        if seed is not None:
            np.random.seed(seed)

        # Sorting state per cube: -1 = unsorted, 0 = wrong bin, 1 = correct bin
        self._cube_status = [-1] * len(self._cubes)
        self._step_count = 0

        # Gripper state: which cube each agent is holding (None = empty hand)
        self._held_cubes = [None] * self.num_agents

        # Collision tracking: which agent pairs were colliding last step
        self._colliding_pairs = set()

        # World reset returns all dynamic prims to their creation poses
        self._world.reset()

        # Slightly randomise cube positions so reset is not too trivial
        for i, info in enumerate(self._cube_info):
            base_pos = info["position"].copy()
            offset = np.array(
                [
                    np.random.uniform(-0.2, 0.2),
                    np.random.uniform(-0.2, 0.2),
                    0.0,
                ]
            )
            self._cubes[i].set_world_pose(
                position=base_pos + offset,
                orientation=np.array([1.0, 0.0, 0.0, 0.0]),
            )

        # Zero robot velocities if the API supports it
        for robot in self._robots:
            try:
                robot.set_linear_velocity(np.zeros(3, dtype=np.float32))
                robot.set_angular_velocity(np.zeros(3, dtype=np.float32))
            except Exception:
                pass

        # Initialise camera annotators (must be done after the stage is reset)
        for camera in self._cameras:
            try:
                camera.initialize()
            except Exception as e:
                print("[WARN] Camera initialisation failed:", e)

        # Start playing and render an initial frame for the cameras
        self._world.play()
        self._world.step(render=True)

        # Initialise potential-based shaping distances
        self._prev_cube_bin_dist = []
        for ci, cube in enumerate(self._cubes):
            cpos, _ = cube.get_world_pose()
            cube_name = self._cube_info[ci]["name"]
            target_bin = None
            for bi in self._bin_info:
                if bi["name"] == cube_name:
                    target_bin = bi
                    break
            if target_bin is not None:
                d = float(np.linalg.norm(cpos[:2] - target_bin["position"][:2]))
            else:
                d = 0.0
            self._prev_cube_bin_dist.append(d)

        self._prev_robot_cube_dist = []
        for ai in range(self.num_agents):
            rpos, _ = self._robots[ai].get_world_pose()
            best_d = float("inf")
            for ci, cube in enumerate(self._cubes):
                cpos, _ = cube.get_world_pose()
                d = np.linalg.norm(rpos[:2] - cpos[:2])
                if d < best_d:
                    best_d = d
            self._prev_robot_cube_dist.append(best_d if best_d != float("inf") else 0.0)

        # Save stage if requested
        if os.environ.get("SAVE_STAGE"):
            self._save_stage()

        return self.get_observations()

    def step(self, actions):
        """
        Advance one simulation step.
        actions: np.ndarray of shape (num_agents, 4) with
                 [forward, lateral, angular, gripper] in the range [-1, 1].
                 gripper > 0.5 activates the sticky gripper.
        Returns: observations, rewards, done, info
        """
        actions = np.asarray(actions, dtype=np.float32)
        if actions.shape == (self.num_agents, 3):
            # Backward-compatible: pad missing gripper action with 0 (open)
            actions = np.hstack(
                [actions, np.zeros((self.num_agents, 1), dtype=np.float32)]
            )
        if actions.shape != (self.num_agents, 4):
            raise ValueError(
                f"actions must have shape (num_agents, 4), got {actions.shape}"
            )

        # Scale and clip
        actions = np.clip(actions, -1.0, 1.0)

        for i in range(self.num_agents):
            forward = actions[i, 0] * self.max_linear_speed
            lateral = actions[i, 1] * self.max_linear_speed
            angular = actions[i, 2] * self.max_angular_speed
            gripper = actions[i, 3]

            if self._robot_type == "kaya":
                # Holonomic Kaya: forward, lateral and yaw are all meaningful
                cmd = np.array([forward, lateral, angular], dtype=np.float32)
                action = self._controllers[i].forward(cmd)
                self._robots[i].apply_wheel_actions(action)

            elif self._robot_type == "jetbot":
                # Differential Jetbot: lateral is not physically possible, so it is
                # ignored (with a one-time warning if non-zero).
                if abs(lateral) > 1e-3:
                    print(
                        "[WARN] Jetbot is a differential-drive robot; "
                        f"lateral action {lateral:.3f} for agent {i} is ignored."
                    )
                cmd = np.array([forward, angular], dtype=np.float32)
                action = self._controllers[i].forward(cmd)
                self._robots[i].apply_wheel_actions(action)

            else:
                # Primitive box: directly set world-frame velocities from body-frame
                # command, giving full 3-DOF motion.
                robot = self._robots[i]
                pos, quat = robot.get_world_pose()
                theta = _yaw_from_quat(quat)
                ct, st = math.cos(theta), math.sin(theta)
                vx_w = ct * forward - st * lateral
                vy_w = st * forward + ct * lateral
                robot.set_linear_velocity(
                    np.array([vx_w, vy_w, 0.0], dtype=np.float32)
                )
                robot.set_angular_velocity(
                    np.array([0.0, 0.0, angular], dtype=np.float32)
                )

            # ------------------------------------------------------------------
            # Simplified sticky gripper: grab / release cubes
            # ------------------------------------------------------------------
            self._update_gripper(i, gripper)

        # Step one physics + render frame
        self._world.step(render=True)
        self._step_count += 1

        # Keep held cubes attached to their robot after the physics step
        self._maintain_held_cubes()

        observations = self.get_observations()
        rewards = self.get_rewards()
        done = False
        info = {"step": self._step_count}

        return observations, rewards, done, info

    # --------------------------------------------------------------------------
    # Gripper helpers
    # --------------------------------------------------------------------------
    def _update_gripper(self, agent_idx, gripper_action):
        """Grab or release the nearest cube based on the gripper action."""
        if gripper_action > self._gripper_threshold:
            # Try to grab if not already holding something
            if self._held_cubes[agent_idx] is not None:
                return
            robot_pos, _ = self._robots[agent_idx].get_world_pose()
            best_cube = None
            best_dist = float("inf")
            for ci, cube in enumerate(self._cubes):
                if self._cube_status[ci] == 1:
                    continue  # already sorted
                if ci in self._held_cubes:
                    continue  # held by another agent
                cpos, _ = cube.get_world_pose()
                dist = np.linalg.norm(robot_pos[:2] - cpos[:2])
                if dist < self._grab_distance and dist < best_dist:
                    best_dist = dist
                    best_cube = ci
            if best_cube is not None:
                self._held_cubes[agent_idx] = best_cube
        else:
            # Release
            self._held_cubes[agent_idx] = None

    def _maintain_held_cubes(self):
        """Position held cubes directly in front of the holding robot."""
        for agent_idx, cube_idx in enumerate(self._held_cubes):
            if cube_idx is None:
                continue
            robot_pos, robot_quat = self._robots[agent_idx].get_world_pose()
            theta = _yaw_from_quat(robot_quat)
            # Hold cube 0.25 m in front of the robot
            offset = 0.25
            cx = robot_pos[0] + math.cos(theta) * offset
            cy = robot_pos[1] + math.sin(theta) * offset
            cz = 0.15  # lift the cube slightly
            self._cubes[cube_idx].set_world_pose(
                position=np.array([cx, cy, cz], dtype=np.float32),
                orientation=np.array([1.0, 0.0, 0.0, 0.0]),
            )
            # Zero its velocity so physics doesn't fight us
            try:
                self._cubes[cube_idx].set_linear_velocity(
                    np.zeros(3, dtype=np.float32)
                )
                self._cubes[cube_idx].set_angular_velocity(
                    np.zeros(3, dtype=np.float32)
                )
            except Exception:
                pass

    def get_observations(self):
        """
        Return a dict with per-agent RGB images and low-dim states.
        rgb:   np.ndarray, shape (num_agents, 128, 128, 3), uint8
        state: np.ndarray, shape (num_agents, 12) =
               [x, y, yaw, vx, vy, omega,
                nearest_cube_dx, nearest_cube_dy, nearest_cube_dist,
                target_bin_dx, target_bin_dy, target_bin_dist]
        """
        rgb = np.zeros((self.num_agents, 128, 128, 3), dtype=np.uint8)
        state = np.zeros((self.num_agents, 12), dtype=np.float32)

        for i, (robot, camera) in enumerate(zip(self._robots, self._cameras)):
            # Camera image
            try:
                img = np.asarray(camera.get_rgba())
                if img is not None and img.ndim == 3 and img.shape[2] >= 3:
                    img = img[:, :, :3]
                    # Handle float [0,1] or uint8 data
                    if img.dtype in (np.float32, np.float64):
                        img = (np.clip(img, 0.0, 1.0) * 255.0).astype(np.uint8)
                    else:
                        img = img.astype(np.uint8)
                    # Trim/pad to exactly 128x128
                    h, w = img.shape[:2]
                    if h >= 128 and w >= 128:
                        img = img[h // 2 - 64 : h // 2 + 64, w // 2 - 64 : w // 2 + 64]
                    else:
                        tmp = np.zeros((128, 128, 3), dtype=np.uint8)
                        tmp[:h, :w] = img[:h, :w]
                        img = tmp
                    rgb[i] = img[:128, :128]
            except Exception as e:
                print(f"[WARN] Could not read camera {i}:", e)

            # Low-dim state
            pos, quat = robot.get_world_pose()
            theta = _yaw_from_quat(quat)
            lin_vel = robot.get_linear_velocity()
            ang_vel = robot.get_angular_velocity()

            # Body-frame velocities
            vx_w, vy_w = lin_vel[0], lin_vel[1]
            ct, st = math.cos(theta), math.sin(theta)
            vx_b = ct * vx_w + st * vy_w
            vy_b = -st * vx_w + ct * vy_w
            omega = ang_vel[2]

            # Find nearest unsorted cube and its target bin
            nc_dx = nc_dy = nc_dist = 0.0
            tgt_dx = tgt_dy = tgt_dist = 0.0

            unsorted = [
                ci for ci in range(len(self._cubes))
                if self._cube_status[ci] != 1
            ]
            if unsorted:
                best_ci = None
                best_dist = float("inf")
                for ci in unsorted:
                    cpos, _ = self._cubes[ci].get_world_pose()
                    d = np.linalg.norm(pos[:2] - cpos[:2])
                    if d < best_dist:
                        best_dist = d
                        best_ci = ci
                if best_ci is not None:
                    cpos, _ = self._cubes[best_ci].get_world_pose()
                    nc_dx = float(cpos[0] - pos[0])
                    nc_dy = float(cpos[1] - pos[1])
                    nc_dist = float(best_dist)

                    # Target bin for this cube
                    cube_name = self._cube_info[best_ci]["name"]
                    target_bin = None
                    for bi in self._bin_info:
                        if bi["name"] == cube_name:
                            target_bin = bi
                            break
                    if target_bin is not None:
                        bpos = target_bin["position"]
                        tgt_dx = float(bpos[0] - pos[0])
                        tgt_dy = float(bpos[1] - pos[1])
                        tgt_dist = float(np.linalg.norm(bpos[:2] - pos[:2]))

            state[i] = np.array(
                [
                    pos[0], pos[1], theta, vx_b, vy_b, omega,
                    nc_dx, nc_dy, nc_dist,
                    tgt_dx, tgt_dy, tgt_dist,
                ],
                dtype=np.float32,
            )

        return {"rgb": rgb, "state": state}

    def get_rewards(self):
        """
        Compute per-agent rewards.
        +10  for a newly-correctly-sorted cube
        -5   for a cube placed in the wrong bin
        -1   one-time for robot-robot collision start / out-of-bounds
        -0.1 step penalty (also while colliding)
        + potential-based shaping:
            * robot→nearest-cube approach  (approach_scale * Δdist)
            * cube→target-bin approach     (shaping_scale * Δdist)
          Using F = gamma * phi(s') - phi(s) with phi = -scale * distance.
        """
        rewards = -0.1 * np.ones(self.num_agents, dtype=np.float32)

        # ------------------------------------------------------------------
        # Robot-robot collisions and out-of-bounds
        # Collision: -1.0 one-time when a new collision starts, -0.1 per step
        # while colliding (avoids catastrophic -200 accumulation).
        # ------------------------------------------------------------------
        robot_positions = []
        colliding_pairs = set()
        for i in range(self.num_agents):
            pos_i, _ = self._robots[i].get_world_pose()
            robot_positions.append(pos_i)
            for j in range(i):
                dist = np.linalg.norm(pos_i[:2] - robot_positions[j][:2])
                if dist < 0.25:
                    colliding_pairs.add((j, i))
                    # One-time penalty when collision starts
                    if (j, i) not in self._colliding_pairs:
                        rewards[i] -= 1.0
                        rewards[j] -= 1.0
                    else:
                        # Smaller per-step penalty while still colliding
                        rewards[i] -= 0.1
                        rewards[j] -= 0.1
            if abs(pos_i[0]) > self.floor_boundary or abs(pos_i[1]) > self.floor_boundary:
                rewards[i] -= 1.0

        self._colliding_pairs = colliding_pairs

        # ------------------------------------------------------------------
        # Cube sorting + potential-based shaping
        # ------------------------------------------------------------------
        for ci, cube in enumerate(self._cubes):
            if self._cube_status[ci] == 1:
                continue  # already correctly sorted

            cpos, _ = cube.get_world_pose()
            cube_name = self._cube_info[ci]["name"]

            # Find the matching (correct) bin and the closest bin
            target_bin = None
            for bi in self._bin_info:
                if bi["name"] == cube_name:
                    target_bin = bi
                    break
            if target_bin is None:
                continue

            target_pos = target_bin["position"]
            curr_cube_to_bin = float(np.linalg.norm(cpos[:2] - target_pos[:2]))

            # Potential-based shaping: cube → target bin
            prev_cube_to_bin = self._prev_cube_bin_dist[ci]
            # F = gamma * phi(s') - phi(s), phi = -scale * dist
            shaping = self._gamma * (-self._shaping_scale * curr_cube_to_bin) - (
                -self._shaping_scale * prev_cube_to_bin
            )
            # Assign shaping to the agent holding this cube, or the nearest agent
            holder = None
            for ai, held in enumerate(self._held_cubes):
                if held == ci:
                    holder = ai
                    break
            if holder is not None:
                rewards[holder] += shaping
            else:
                # Assign to nearest agent to the cube
                best_ai = 0
                best_d = float("inf")
                for ai in range(self.num_agents):
                    d = np.linalg.norm(robot_positions[ai][:2] - cpos[:2])
                    if d < best_d:
                        best_d = d
                        best_ai = ai
                rewards[best_ai] += shaping * 0.5  # reduced credit when not holding

            self._prev_cube_bin_dist[ci] = curr_cube_to_bin

            # Check if cube is in a bin
            closest_idx = None
            closest_dist = float("inf")
            for bi, bin_info in enumerate(self._bin_info):
                dist = np.linalg.norm(cpos[:2] - bin_info["position"][:2])
                if dist < closest_dist:
                    closest_dist = dist
                    closest_idx = bi

            if closest_dist < self._bin_threshold and closest_idx is not None:
                bin_name = self._bin_info[closest_idx]["name"]
                if bin_name == cube_name:
                    if self._cube_status[ci] != 1:
                        rewards += 10.0
                        self._cube_status[ci] = 1
                        print(
                            f"[ENV] Cube {ci} ({cube_name}) correctly sorted! "
                            f"+10 reward to all agents."
                        )
                else:
                    if self._cube_status[ci] == -1:
                        rewards -= 5.0
                        self._cube_status[ci] = 0
                        print(
                            f"[ENV] Cube {ci} ({cube_name}) in wrong bin "
                            f"({bin_name}). -5 reward."
                        )

        # ------------------------------------------------------------------
        # Robot→nearest-cube approach shaping (encourages exploration)
        # ------------------------------------------------------------------
        for ai in range(self.num_agents):
            rpos = robot_positions[ai]
            unsorted = [
                ci for ci in range(len(self._cubes))
                if self._cube_status[ci] != 1
            ]
            if not unsorted:
                continue
            best_ci = None
            best_d = float("inf")
            for ci in unsorted:
                cpos, _ = self._cubes[ci].get_world_pose()
                d = np.linalg.norm(rpos[:2] - cpos[:2])
                if d < best_d:
                    best_d = d
                    best_ci = ci
            if best_ci is not None:
                prev_d = self._prev_robot_cube_dist[ai]
                # F = gamma * phi(s') - phi(s), phi = -scale * dist
                approach = self._gamma * (-self._approach_scale * best_d) - (
                    -self._approach_scale * prev_d
                )
                rewards[ai] += approach
                self._prev_robot_cube_dist[ai] = best_d

        return rewards

    def render(self):
        """Render a fresh frame and return the current observations.
        Note: in headless standalone mode a render step is a regular simulation
        step, so this advances physics by one frame.
        """
        self._world.step(render=True)
        return self.get_observations()

    def close(self):
        """Stop simulation and close the application."""
        try:
            self._world.stop()
        except Exception:
            pass
        SIMULATION_APP.close()

    # --------------------------------------------------------------------------
    # Internal helpers
    # --------------------------------------------------------------------------
    def _save_stage(self):
        """Save the current USD stage to OUTPUT_DIR if SAVE_STAGE is set."""
        try:
            import omni.usd

            os.makedirs(OUTPUT_DIR, exist_ok=True)
            timestamp = time.strftime("%Y%m%d_%H%M%S")
            path = os.path.join(
                OUTPUT_DIR, f"multi_agent_stage_{timestamp}.usd"
            ).replace("\\", "/")
            stage = omni.usd.get_context().get_stage()
            # Export the stage (flattened) without opening a new context
            ok = stage.Export(path) if stage else False
            print(f"[INFO] Saved stage to {path} (ok={ok})")
        except Exception as e:
            print("[WARN] Could not save stage:", e)


# -----------------------------------------------------------------------------
# Standalone smoke test
# -----------------------------------------------------------------------------
def main():
    print("=" * 60)
    print("Multi-Agent Isaac Sim Environment - Smoke Test")
    print("=" * 60)

    num_agents = int(os.environ.get("NUM_AGENTS", NUM_AGENTS))
    env = MultiAgentIsaacEnv(num_agents=num_agents, try_assets=True)

    try:
        obs = env.reset()
        print("Initial observations:")
        print("  rgb  shape:", obs["rgb"].shape)
        print("  state shape:", obs["state"].shape)

        for t in range(10):
            actions = np.random.uniform(-0.5, 0.5, size=(num_agents, 4))
            obs, rewards, done, info = env.step(actions)
            print(
                f"step {t:02d}: rewards={rewards}, "
                f"agent0_state=[{obs['state'][0][0]:.2f}, "
                f"{obs['state'][0][1]:.2f}, "
                f"{obs['state'][0][2]:.2f}, "
                f"nc_dist={obs['state'][0][8]:.2f}, "
                f"tgt_dist={obs['state'][0][11]:.2f}]"
            )

        # Save an example frame for visual sanity-check
        try:
            os.makedirs(OUTPUT_DIR, exist_ok=True)
            from PIL import Image

            sample_img = obs["rgb"][0]
            Image.fromarray(sample_img).save(
                os.path.join(OUTPUT_DIR, "smoke_agent0_rgb.png")
            )
            print("Saved example image to", OUTPUT_DIR)
        except Exception as e:
            print("[INFO] Could not save example image:", e)

    except Exception as e:
        print("[ERROR] Smoke test failed:")
        traceback.print_exc()
    finally:
        try:
            env.close()
        except Exception:
            pass
        print("Smoke test finished.")
        # Skip Py_FinalizeEx to avoid a known 5.1.0/D3D12 shutdown crash.
        import sys
        sys.stdout.flush()
        sys.stderr.flush()
        os._exit(0)


if __name__ == "__main__":
    main()