OmniverseのUSDシーンを生成AIと連携して操作する

# 1. はじめに
NVIDIAのOmniverseの拡張機能を作成し、自然言語でOmniverseのUSDシーンを操作できるようにしてみました。自然言語を入力すると、生成AIがUSDシーンを操作するPythonコードを生成します。NVIDIA謹製の[USD Code API](https://docs.omniverse.nvidia.com/services/latest/services/usd-code/overview.html)と拡張機能のChat USDで自然言語でのシーン操作がサポートされていますが、今回は自分で作ってみました。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/4072158/0cf3b9ea-1ff7-4f05-b3db-4986a71eb3ce.png)

# 2. どのくらいのものができたか
オブジェクトの追加、移動、回転の付与などの基本操作ができる程度のものはできました。

10個のCubeを配置。
![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/4072158/9270d67f-ded4-4152-ab26-16913f6b9ba6.png)

いくつかのCubeを回転。
![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/4072158/f67bd65a-9f96-4a53-952b-eb9ddec49e72.png)

2つめのCubeを移動。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/4072158/ff6ecb88-179c-48b5-a358-7cf224f641d5.png)

2つめのCubeをもっと移動。
![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/4072158/3db40f24-937d-4767-9d2e-3b23cf6d10d4.png)

# 3. 作り方(ポイントのみ)
ポイントのみ解説します。コードの全文は最後に載せています。

## 3.1 プロンプト生成部分
シーン情報を取得する部分がポイントです。USDA形式でシーンの情報を取得し、これをプロンプトに含めています。こうすることで生成AIが現在のUSDシーンがどうなっているかを考慮したコードを生成することができるようになります。

```python
    # シーン情報をUSD（Universal Scene Description）形式の文字列
    def _get_scene_info(self) -> str:
        stage = omni.usd.get_context().get_stage()
        return stage.GetRootLayer().ExportToString() if stage else ""

    def _generate_ai_prompt(self, user_input: str) -> str:
        return f"""あなたはPythonを使ったUSDシーン操作のエキスパートです。
                ユーザーの記述に基づいてUSDシーンを修正するPythonコードを生成してください。
                pxrライブラリ, omniライブラリを使用してください。
                実行可能なPythonのコードブロックのみを返し、説明やMarkdown形式は含めないでください。
                Stageの取得は次のコードを使用してください。

                # Stageの取得に用いるコード
                stage = omni.usd.get_context().get_stage()

                なお、現在のUSDシーンの状態は以下です。

                # USDシーンの状態
                {self._get_scene_info()}"""
```

## 3.2 コード実行部分
生成AIが生成したPythonコードをexec関数に引き渡しているだけ、特に難しいことはしてません。

```python
    def _execute_python_code(self, code: str) -> None:
        if not code:
            return

        try:
            exec(code)
            self._code_execution_error_info.set_value("Success!!")
        except Exception as e:
            carb.log_error(f"Error executing code: {str(e)}")
            self._code_execution_error_info.set_value(str(e))

```

# 4. 課題
(1) CONTEXT LENGTHの最大長の制限にひっかかる。
USDシーンの情報が膨大になってきたときに、生成AIのCONTEXT LENGTHの制限にひっかかってコード生成ができません。 コード生成に必要なUSDシーンの情報だけに絞って生成AIに渡す工夫が必要そうです。

(2) 動作しないコードが生成される。
しばしば動作しないコードが生成されます。GPT-4oとDeepSeek Chatの2パターンを試しましたが、どちらでも似たようなものでした。プロンプトをもっと工夫するか、RAGの利用、もしくはファインチューニングなど、生成されるコードの品質を高める工夫が必要と感じました。

# 5. まとめ
今回は、生成AIと連携して、自然言語でOmniverseのUSDシーンを操作する拡張機能を作成してみました。最低限の動作をするものは作れましたが、改善の余地が残りました。

# 6. コード全文
拡張機能のコード全文は下記となります。
生成AI使いつつ、お試しに作っただけなので、いろいろ変なところあるかもしれません。

<details><summary>コード全文</summary>

```python
import omni.ui as ui
import omni.kit.ui
import omni.kit.commands
from omni import usd
from pxr import Sdf
import carb
import json
import requests
import re
import asyncio

class MyExtension(omni.ext.IExt):
    API_KEY = "APIキー"
    API_URL = "https://api.deepseek.com/v1/chat/completions"
    MODEL_NAME = "deepseek-coder"

    def on_startup(self, ext_id):
        self._window = None
        self._user_input = None
        self._generated_code = None
        self._code_execution_error_info = None
        self._generating_label = None
        self._create_window()

    def on_shutdown(self):
        if self._window:
            self._window.destroy()
            self._window = None


    def _toggle_window_visibility(self, menu_path: str, visible: bool):
        if visible:
            self._create_window()
        elif self._window:
            self._window = None

    def _create_window(self):
        self._window = ui.Window("AI USD Scene Modifier", width=500, height=600)
        with self._window.frame:
            with ui.VStack():
                # Input Section
                ui.Label("Describe your scene changes:")
                self._user_input = ui.StringField(height=100, multiline=True).model

                # Generation Section
                ui.Button("Generate Python Code", clicked_fn=self._on_generate_clicked, height=30)
                ui.Label("Generated Python Code:")
                self._generated_code = ui.StringField(height=200, multiline=True, read_only=False).model

                # Execution Section
                ui.Button("Execute Code", clicked_fn=self._on_execute_clicked, height=30)
                ui.Label("Code execution result:")
                self._code_execution_error_info = ui.StringField(height=100, multiline=True, read_only=True).model

        self._window.set_visibility_changed_fn(self._visiblity_changed_fn)

    def _get_scene_info(self) -> str:
        stage = omni.usd.get_context().get_stage()
        return stage.GetRootLayer().ExportToString() if stage else ""

    def _execute_python_code(self, code: str) -> None:
        if not code:
            return

        try:
            exec(code)
            self._code_execution_error_info.set_value("Success!!")
        except Exception as e:
            carb.log_error(f"Error executing code: {str(e)}")
            self._code_execution_error_info.set_value(str(e))

    def _on_execute_clicked(self):
        code = self._generated_code.get_value_as_string()
        self._execute_python_code(code)

    def _visiblity_changed_fn(self, visible: bool):
        editor_menu = omni.kit.ui.get_editor_menu()
        if editor_menu:
            editor_menu.set_value("Window/My Company Window", visible)

    def _extract_python_code(self, markdown_text: str) -> str:
        match = re.search(r"```python\n(.*?)\n```", markdown_text, re.DOTALL)
        return match.group(1) if match else markdown_text

    def _generate_ai_prompt(self, user_input: str) -> str:
        return f"""あなたはPythonを使ったUSDシーン操作のエキスパートです。
                ユーザーの記述に基づいてUSDシーンを修正するPythonコードを生成してください。
                pxrライブラリ, omniライブラリを使用してください。
                実行可能なPythonのコードブロックのみを返し、説明やMarkdown形式は含めないでください。
                Stageの取得は次のコードを使用してください。

                # Stageの取得に用いるコード
                stage = omni.usd.get_context().get_stage()

                なお、現在のUSDシーンの状態は以下です。

                # USDシーンの状態
                {self._get_scene_info()}"""



    def _call_ai_api(self, prompt: str, user_input: str) -> str:
        headers = {
            "Content-Type": "application/json",
            "Authorization": f"Bearer {self.API_KEY}"
        }

        payload = {
            "model": self.MODEL_NAME,
            "messages": [
                {"role": "system", "content": prompt},
                {"role": "user", "content": user_input}
            ]
        }
        print("Payload:", json.dumps(payload, indent=2))

        response = requests.post(self.API_URL, headers=headers, json=payload)
        response.raise_for_status()
        return response.json()["choices"][0]["message"]["content"]

    def _on_generate_clicked(self):
        self._generated_code.set_value("")
        user_input = self._user_input.get_value_as_string()

        try:
            prompt = self._generate_ai_prompt(user_input)
            ai_response = self._call_ai_api(prompt, user_input)
            code = self._extract_python_code(ai_response)
            self._generated_code.set_value(code)
        except Exception as e:
            carb.log_error(f"Generative AI API error: {str(e)}")
            self._generated_code.set_value(f"API Error: {str(e)}")

```　

</details>