demo1.py

# Core Omniverse Kit and USD APIs for scene manipulation and simulation control.
import omni.kit.commands
import omni.usd
import omni.timeline
from pxr import Sdf, UsdGeom, Gf, UsdPhysics

# Cesium integration for real-world geospatial positioning.
# NOTE: We are NOT importing UsdPhysics, as we will move the object kinematically.
from cesium.usd.plugins.CesiumUsdSchemas import GlobeAnchorAPI as CesiumGlobeAnchorAPI

# Standard Python libraries for math and logging.
import math
import logging

# --- Script Configuration ---
DRONE_PATH = "/World/cf2x" # This is using the already existing drone in the scene
INITIAL_SCALE = 100.0
# Define the drone's starting coordinates to reset to on every run.
INITIAL_LATITUDE = 37.2201289
INITIAL_LONGITUDE = -80.4181060
INITIAL_HEIGHT = 635.0  # Height in meters above the stadium
MOVEMENT_DISTANCE = 0.0003
MOVEMENT_SPEED = 0.5

logging.basicConfig(level=logging.INFO, format="%(levelname)s: %(message)s")

# Global state dictionary to hold simulation variables.
g_simulation_state = {}


def on_update_movement(e):
    """
    This function is called by Isaac Sim on every frame of the simulation.
    It calculates and applies the new position for the drone.
    """
    timeline = g_simulation_state.get("timeline")
    if not timeline or not timeline.is_playing():
        return

    globe_anchor = g_simulation_state.get("globe_anchor")
    start_time = g_simulation_state.get("start_time")
    initial_pos = g_simulation_state.get("initial_position")

    if not all([globe_anchor, start_time is not None, initial_pos]):
        return # Silently return if not yet initialized

    # Calculate and apply new position
    elapsed_time = timeline.get_current_time() - start_time
    displacement = math.sin(elapsed_time * MOVEMENT_SPEED) * MOVEMENT_DISTANCE
    new_longitude = initial_pos["longitude"] + displacement
    globe_anchor.GetAnchorLongitudeAttr().Set(new_longitude)


def stop_movement():
    """Stops the drone's movement by cancelling the frame-by-frame update subscription."""
    if g_simulation_state.get("update_subscription"):
        g_simulation_state["update_subscription"].unsubscribe()
        g_simulation_state["update_subscription"] = None
        logging.info("Drone movement stopped.")


def setup_and_move_drone():
    """
    This is the main function. It finds the drone, cleans it of any old physics components,
    sets its scale and anchor, and starts the movement subscription.
    """
    stop_movement()

    context = omni.usd.get_context()
    stage = context.get_stage()
    timeline = omni.timeline.get_timeline_interface()
    app_update = omni.kit.app.get_app().get_update_event_stream()

    drone_prim = stage.GetPrimAtPath(DRONE_PATH)
    if not drone_prim.IsValid():
        logging.error(f"Drone not found at path '{DRONE_PATH}'.")
        return

    # --- THIS IS THE CRITICAL FIX ---
    # Before doing anything, check if a stale RigidBodyAPI from a previous run
    # exists on the prim. If so, remove it to prevent conflicts.
    if drone_prim.HasAPI(UsdPhysics.RigidBodyAPI):
        logging.warning("Found stale RigidBodyAPI on drone. Removing it to ensure a clean state.")
        drone_prim.RemoveAPI(UsdPhysics.RigidBodyAPI)
    # ---------------------------------

    # Set scale (this is a safe, non-physics operation).
    xform_api = UsdGeom.XformCommonAPI(drone_prim)
    xform_api.SetScale(Gf.Vec3f(INITIAL_SCALE, INITIAL_SCALE, INITIAL_SCALE))
    logging.info(f"Drone scale set to {INITIAL_SCALE}x.")

    # Apply Cesium Globe Anchor. This is all we need to control its position.
    globe_anchor = CesiumGlobeAnchorAPI.Apply(drone_prim)
    if not globe_anchor:
        logging.error("Failed to apply CesiumGlobeAnchorAPI to the drone prim.")
        return

    # --- POSITION RESET ---
    # Explicitly set the drone's position to the configured initial coordinates.
    # This acts as a "reset" every time the script is run.
    globe_anchor.GetAnchorLatitudeAttr().Set(INITIAL_LATITUDE)
    globe_anchor.GetAnchorLongitudeAttr().Set(INITIAL_LONGITUDE)
    globe_anchor.GetAnchorHeightAttr().Set(INITIAL_HEIGHT)
    logging.info(f"Drone position reset to: Lat={INITIAL_LATITUDE}, Lon={INITIAL_LONGITUDE}, H={INITIAL_HEIGHT}")

    # Store the initial state and start the update loop.
    g_simulation_state.update({
        "timeline": timeline,
        "globe_anchor": globe_anchor,
        "start_time": timeline.get_current_time(),
        "initial_position": {
            "latitude": globe_anchor.GetAnchorLatitudeAttr().Get(),
            "longitude": globe_anchor.GetAnchorLongitudeAttr().Get(),
            "height": globe_anchor.GetAnchorHeightAttr().Get(),
        },
        "update_subscription": app_update.create_subscription_to_pop(on_update_movement)
    })

    logging.info("Drone controller is active. Press PLAY to see movement.")
    print("--- To stop the movement, run the command: stop_movement() ---")


# --- Main Execution ---
setup_and_move_drone()