spawn_demo.py

# Copyright (c) 2022-2024, The Isaac Lab Project Developers.
# All rights reserved.
#
# SPDX-License-Identifier: BSD-3-Clause

import argparse
import numpy as np
import torch
from isaaclab.app import AppLauncher


# add argparse arguments
parser = argparse.ArgumentParser(
    description="This script demonstrates different single-arm manipulators."
)
# append AppLauncher cli args
AppLauncher.add_app_launcher_args(parser)
# parse the arguments
args_cli = parser.parse_args()

# launch omniverse app
app_launcher = AppLauncher(args_cli)
simulation_app = app_launcher.app

# isaac sim modules are only available after launching it (?)

import carb
import gymnasium as gym

from pxr import Usd

import isaacsim.core.utils.prims as prim_utils

import isaaclab.sim as sim_utils
from isaaclab.assets import Articulation
from isaaclab.utils.assets import ISAAC_NUCLEUS_DIR, ISAACLAB_NUCLEUS_DIR

from simpub.sim.isaacsim_publisher import IsaacSimPublisher

print(ISAACLAB_NUCLEUS_DIR)

##
# Pre-defined configs
##
# isort: off
from isaaclab_assets import (
    FRANKA_PANDA_CFG,
    UR10_CFG,
    KINOVA_JACO2_N7S300_CFG,
    KINOVA_JACO2_N6S300_CFG,
    KINOVA_GEN3_N7_CFG,
    SAWYER_CFG,
)

# isort: on


def define_origins(num_origins: int, spacing: float) -> list[list[float]]:
    """Defines the origins of the the scene."""
    # currently this function is unimportant, since we only test with a single origin/env.

    # create tensor based on number of environments
    env_origins = torch.zeros(num_origins, 3)
    # create a grid of origins
    num_rows = np.floor(np.sqrt(num_origins))
    num_cols = np.ceil(num_origins / num_rows)
    xx, yy = torch.meshgrid(
        torch.arange(num_rows), torch.arange(num_cols), indexing="xy"
    )
    env_origins[:, 0] = (
        spacing * xx.flatten()[:num_origins] - spacing * (num_rows - 1) / 2
    )
    env_origins[:, 1] = (
        spacing * yy.flatten()[:num_origins] - spacing * (num_cols - 1) / 2
    )
    env_origins[:, 2] = 0.0
    # return the origins
    return env_origins.tolist()


def design_scene() -> tuple[dict, list[list[float]]]:
    """Designs the scene."""
    # this function will build the scene by adding primitives to it.

    # Ground-plane
    cfg = sim_utils.GroundPlaneCfg()
    cfg.func("/World/defaultGroundPlane", cfg)

    # Lights
    cfg = sim_utils.DomeLightCfg(intensity=2000.0, color=(0.75, 0.75, 0.75))
    cfg.func("/World/Light", cfg)

    # Create separate groups called "Origin1", "Origin2", "Origin3"
    # Each group will have a mount and a robot on top of it
    origins = define_origins(num_origins=1, spacing=2.0)

    # Origin 1 with Franka Panda
    prim_utils.create_prim("/World/Origin1", "Xform", translation=origins[0])
    prim_utils.create_prim("/World/Origin1/Tables", "Xform")

    cfg_wall_1 = sim_utils.CuboidCfg(
        size=(10, 0.1, 1),
        collision_props=sim_utils.CollisionPropertiesCfg(),
        visual_material=sim_utils.PreviewSurfaceCfg(
            diffuse_color=(1.0, 0.7, 0.4)
        ),
    )
    cfg_wall_1.func("/World/Wall1", cfg_wall_1, translation=(0, -5, 0))
    cfg_wall_1.func("/World/Wall2", cfg_wall_1, translation=(0, 5, 0))

    cfg_wall_2 = sim_utils.CuboidCfg(
        size=(0.1, 10, 1),
        collision_props=sim_utils.CollisionPropertiesCfg(),
        visual_material=sim_utils.PreviewSurfaceCfg(
            diffuse_color=(1.0, 1.0, 1.0)
        ),
    )
    cfg_wall_2.func("/World/Wall3", cfg_wall_2, translation=(-5, 0, 0))
    cfg_wall_2.func("/World/Wall4", cfg_wall_2, translation=(5, 0, 0))

    # # -- Table
    cfg = sim_utils.UsdFileCfg(
        usd_path=f"{ISAAC_NUCLEUS_DIR}/Props/Mounts/SeattleLabTable/table_instanceable.usd"
    )
    cfg.func("/World/Origin1/Tables/Table", cfg, translation=(0.55, 0.0, 1.05))
    # cfg.func("/World/Origin1/Tables/Table_1", cfg, translation=(0.55, 3.0, 1.05))

    # -- Robot
    franka_arm_cfg = FRANKA_PANDA_CFG.replace(prim_path="/World/Origin1/Robot")
    franka_arm_cfg.init_state.pos = (0.0, 0.0, 1.05)
    franka_panda = Articulation(cfg=franka_arm_cfg)

    # -- cube
    cfg_cube = sim_utils.CuboidCfg(
        size=(0.1, 0.1, 0.1),
        rigid_props=sim_utils.RigidBodyPropertiesCfg(),
        mass_props=sim_utils.MassPropertiesCfg(mass=1.0),
        collision_props=sim_utils.CollisionPropertiesCfg(),
        visual_material=sim_utils.PreviewSurfaceCfg(
            diffuse_color=(0.0, 1.0, 1.0)
        ),
    )
    cfg_cube.func(
        "/World/Origin1/Cube1", cfg_cube, translation=(0.2, 0.0, 3.0)
    )

    cfg_cube = sim_utils.CuboidCfg(
        size=(0.1, 0.2, 0.3),
        rigid_props=sim_utils.RigidBodyPropertiesCfg(),
        mass_props=sim_utils.MassPropertiesCfg(mass=1.0),
        collision_props=sim_utils.CollisionPropertiesCfg(),
        visual_material=sim_utils.PreviewSurfaceCfg(
            diffuse_color=(0.0, 1.0, 1.0)
        ),
    )
    cfg_cube.func(
        "/World/Origin1/Cube2", cfg_cube, translation=(0.2, 0.0, 5.0)
    )

    cfg_capsule = sim_utils.CapsuleCfg(
        radius=0.1,
        height=0.3,
        axis="Z",
        rigid_props=sim_utils.RigidBodyPropertiesCfg(),
        mass_props=sim_utils.MassPropertiesCfg(mass=1.0),
        collision_props=sim_utils.CollisionPropertiesCfg(),
        visual_material=sim_utils.PreviewSurfaceCfg(
            diffuse_color=(1.0, 0.1, 0.1)
        ),
    )
    cfg_capsule.func(
        "/World/Origin1/CapsuleZ", cfg_capsule, translation=(0.2, 0.0, 7.0)
    )

    cfg_capsule = sim_utils.CapsuleCfg(
        radius=0.2,
        height=0.4,
        axis="X",
        rigid_props=sim_utils.RigidBodyPropertiesCfg(),
        mass_props=sim_utils.MassPropertiesCfg(mass=1.0),
        collision_props=sim_utils.CollisionPropertiesCfg(),
        visual_material=sim_utils.PreviewSurfaceCfg(
            diffuse_color=(0.1, 1.0, 0.1)
        ),
    )
    cfg_capsule.func(
        "/World/Origin1/CapsuleX", cfg_capsule, translation=(0.2, 0.0, 9.0)
    )

    cfg_capsule = sim_utils.CapsuleCfg(
        radius=0.3,
        height=0.5,
        axis="Y",
        rigid_props=sim_utils.RigidBodyPropertiesCfg(),
        mass_props=sim_utils.MassPropertiesCfg(mass=1.0),
        collision_props=sim_utils.CollisionPropertiesCfg(),
        visual_material=sim_utils.PreviewSurfaceCfg(
            diffuse_color=(0.1, 0.1, 1.0)
        ),
    )
    cfg_capsule.func(
        "/World/Origin1/CapsuleY", cfg_capsule, translation=(0.2, 0.0, 11.0)
    )

    cfg_cone = sim_utils.ConeCfg(
        radius=0.1,
        height=0.3,
        axis="Z",
        rigid_props=sim_utils.RigidBodyPropertiesCfg(),
        mass_props=sim_utils.MassPropertiesCfg(mass=1.0),
        collision_props=sim_utils.CollisionPropertiesCfg(),
        visual_material=sim_utils.PreviewSurfaceCfg(
            diffuse_color=(1.0, 0.1, 0.1)
        ),
    )
    cfg_cone.func(
        "/World/Origin1/ConeZ", cfg_cone, translation=(0.2, 0.0, 13.0)
    )

    cfg_cone = sim_utils.ConeCfg(
        radius=0.2,
        height=0.4,
        axis="X",
        rigid_props=sim_utils.RigidBodyPropertiesCfg(),
        mass_props=sim_utils.MassPropertiesCfg(mass=1.0),
        collision_props=sim_utils.CollisionPropertiesCfg(),
        visual_material=sim_utils.PreviewSurfaceCfg(
            diffuse_color=(0.1, 1.0, 0.1)
        ),
    )
    cfg_cone.func(
        "/World/Origin1/ConeX", cfg_cone, translation=(0.2, 0.0, 15.0)
    )

    cfg_cone = sim_utils.ConeCfg(
        radius=0.3,
        height=0.5,
        axis="Y",
        rigid_props=sim_utils.RigidBodyPropertiesCfg(),
        mass_props=sim_utils.MassPropertiesCfg(mass=1.0),
        collision_props=sim_utils.CollisionPropertiesCfg(),
        visual_material=sim_utils.PreviewSurfaceCfg(
            diffuse_color=(0.1, 0.1, 1.0)
        ),
    )
    cfg_cone.func(
        "/World/Origin1/ConeY", cfg_cone, translation=(0.2, 0.0, 17.0)
    )

    cfg_cylinder = sim_utils.CylinderCfg(
        radius=0.1,
        height=0.3,
        axis="Z",
        rigid_props=sim_utils.RigidBodyPropertiesCfg(),
        mass_props=sim_utils.MassPropertiesCfg(mass=1.0),
        collision_props=sim_utils.CollisionPropertiesCfg(),
        visual_material=sim_utils.PreviewSurfaceCfg(
            diffuse_color=(1.0, 0.1, 0.1)
        ),
    )
    cfg_cylinder.func(
        "/World/Origin1/CylinderZ", cfg_cylinder, translation=(0.2, 0.0, 19.0)
    )

    cfg_cylinder = sim_utils.CylinderCfg(
        radius=0.2,
        height=0.4,
        axis="X",
        rigid_props=sim_utils.RigidBodyPropertiesCfg(),
        mass_props=sim_utils.MassPropertiesCfg(mass=1.0),
        collision_props=sim_utils.CollisionPropertiesCfg(),
        visual_material=sim_utils.PreviewSurfaceCfg(
            diffuse_color=(0.1, 1.0, 0.1)
        ),
    )
    cfg_cylinder.func(
        "/World/Origin1/CylinderX", cfg_cylinder, translation=(0.2, 0.0, 21.0)
    )

    cfg_cylinder = sim_utils.CylinderCfg(
        radius=0.3,
        height=0.5,
        axis="Y",
        rigid_props=sim_utils.RigidBodyPropertiesCfg(),
        mass_props=sim_utils.MassPropertiesCfg(mass=1.0),
        collision_props=sim_utils.CollisionPropertiesCfg(),
        visual_material=sim_utils.PreviewSurfaceCfg(
            diffuse_color=(0.1, 0.1, 1.0)
        ),
    )
    cfg_cylinder.func(
        "/World/Origin1/CylinderY", cfg_cylinder, translation=(0.2, 0.0, 23.0)
    )

    cfg_sphere = sim_utils.SphereCfg(
        radius=0.3,
        rigid_props=sim_utils.RigidBodyPropertiesCfg(),
        mass_props=sim_utils.MassPropertiesCfg(mass=1.0),
        collision_props=sim_utils.CollisionPropertiesCfg(),
        visual_material=sim_utils.PreviewSurfaceCfg(
            diffuse_color=(0.1, 0.1, 1.0)
        ),
    )
    cfg_sphere.func(
        "/World/Origin1/Sphere1", cfg_sphere, translation=(0.2, 0.0, 25.0)
    )

    # return the scene information
    scene_entities = {
        "franka_panda": franka_panda,
    }
    return scene_entities, origins


def run_simulator(
    sim: sim_utils.SimulationContext,
    entities: dict[str, Articulation],
    origins: torch.Tensor,
):
    """Runs the simulation loop."""
    # Define simulation stepping
    sim_dt = sim.get_physics_dt()
    sim_time = 0.0
    count = 0
    # Simulate physics
    while simulation_app.is_running():
        # reset
        if count % 200 == 0:
            # reset counters
            sim_time = 0.0
            count = 0
            # reset the scene entities
            for index, robot in enumerate(entities.values()):
                # root state
                root_state = robot.data.default_root_state.clone()
                root_state[:, :3] += origins[index]
                robot.write_root_state_to_sim(root_state)
                # set joint positions
                joint_pos, joint_vel = (
                    robot.data.default_joint_pos.clone(),
                    robot.data.default_joint_vel.clone(),
                )
                robot.write_joint_state_to_sim(joint_pos, joint_vel)
                # clear internal buffers
                robot.reset()
            print("[INFO]: Resetting robots state...")
        # apply random actions to the robots
        for robot in entities.values():
            # generate random joint positions
            joint_pos_target = (
                robot.data.default_joint_pos
                + torch.randn_like(robot.data.joint_pos) * 0.1
            )
            joint_pos_target = joint_pos_target.clamp_(
                robot.data.soft_joint_pos_limits[..., 0],
                robot.data.soft_joint_pos_limits[..., 1],
            )
            # apply action to the robot
            robot.set_joint_position_target(joint_pos_target)
            # write data to sim
            robot.write_data_to_sim()
        # perform step
        sim.step()
        # update sim-time
        sim_time += sim_dt
        count += 1
        # update buffers
        for robot in entities.values():
            robot.update(sim_dt)


def main():
    """Main function."""
    # sim_utils.SimulationContext is a singleton class
    # if SimulationContext.instance() is None:
    #     self.sim: SimulationContext = SimulationContext(self.cfg.sim)
    # or: print(env.sim)

    # Initialize the simulation context
    sim_cfg = sim_utils.SimulationCfg()

    # # use cpu simulation
    # sim_cfg.use_fabric = False
    # sim_cfg.device = "cpu"
    # sim_cfg.use_gpu_pipeline = False
    # sim_cfg.physx.use_gpu = False

    sim = sim_utils.SimulationContext(sim_cfg)

    # Set main camera
    sim.set_camera_view([3.5, 0.0, 3.2], [0.0, 0.0, 0.5])

    # design scene
    scene_entities, scene_origins = design_scene()
    scene_origins = torch.tensor(scene_origins, device=sim.device)

    # Play the simulator
    sim.reset()

    # Now we are ready!
    print("[INFO]: Setup complete...")

    # start publisher
    # make sure you have a correct host address
    _ = IsaacSimPublisher(host="192.168.0.134", stage=sim.stage)

    # Run the simulator
    run_simulator(sim, scene_entities, scene_origins)


if __name__ == "__main__":
    main()
    simulation_app.close()