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()