Add LLMAgent to EnvGraphSpec generation

[qianl-nv]

Summary

First part of the LLM env gen pipeline: LLMAgent going from prompt to LLMEnvSchema (JSON)

Detailed description

• Reason: LLM infers task, robot & object names from prompt
  Changed: Add new LLMAgent class & output shema (defines subset of supported task/constraints ect).
• Impact: Supports the first step in the user journey for agentic env gen

To run: python -m isaaclab_arena.llm_env_gen.try_schema

[isaaclab-review-bot]

Review Summary

Thanks for adding the LLM-based environment generation pipeline! The modular design with clear separation between schema definition, LLM interaction, and deterministic resolution is well thought out. The trace-based debugging will be valuable for understanding resolution decisions.

Suggestions

1. Replace assertions with explicit exceptions for runtime validation

File: isaaclab_arena/llm_env_gen/llm_agent.py (lines 57-58)

The API key validation uses assert, which can be stripped with Python's -O flag:

assert self.api_key, "API key required: set NV_API_KEY or pass api_key."

Consider using an explicit exception:

if not self.api_key:
    raise ValueError("API key required: set NV_API_KEY or pass api_key.")

2. Improve JSON extraction error handling

File: isaaclab_arena/llm_env_gen/llm_agent.py (lines 148-161)

The _extract_json method uses contextlib.suppress(json.JSONDecodeError) before attempting manual parsing. If both paths fail, it raises AssertionError. Consider:

1. Raising a more descriptive custom exception (e.g., LLMResponseParseError)
2. Including the raw response in error messages for debugging

class LLMResponseParseError(ValueError):
    """Raised when the LLM response cannot be parsed as JSON."""

@staticmethod
def _extract_json(content: str) -> dict:
    content = content.strip()
    # ... extraction logic ...
    raise LLMResponseParseError(f"Could not extract JSON from LLM response: {content[:200]}...")

3. Consider adding retry logic for transient API failures

File: isaaclab_arena/llm_env_gen/llm_agent.py

LLM API calls can fail transiently due to rate limits or network issues. Consider adding retry logic with exponential backoff, or at minimum documenting that callers should implement retries:

from tenacity import retry, stop_after_attempt, wait_exponential

@retry(stop=stop_after_attempt(3), wait=wait_exponential(multiplier=1, min=2, max=10))
def generate_spec(self, ...):
    ...

4. Add test coverage for LLMAgent

File: isaaclab_arena/tests/test_resolver.py

The resolver tests are thorough with the FakeAssetRegistry pattern. Consider adding similar tests for LLMAgent:

• Test _extract_json with various edge cases (markdown fences, nested JSON, malformed responses)
• Test schema validation rejections
• Mock the OpenAI client to test the overall flow without making real API calls

5. Document the silent failure behavior in resolver

File: isaaclab_arena/llm_env_gen/resolver.py

The design choice to never raise on LLM mistakes (recording trace events instead) is documented in the class docstring, which is good. Consider adding a helper method or property to check if resolution was "clean":

@property
def has_resolution_errors(self) -> bool:
    """Returns True if any trace events indicate failed resolution."""
    error_stages = {"item.miss", "name.miss", "task.unknown_subject", "task.unknown_target"}
    return any(e.stage in error_stages for e in self.trace)

Minor Notes

• The try_schema.py script is useful for manual testing but consider adding it to .gitignore or documenting it as a development-only tool
• Good use of __post_init__ in ArenaEnvGraphSpec to centralize invariant checking

Overall this is a solid foundation for the LLM environment generation feature.

---

Update (commit 3336f26): The PR has been substantially refactored since my initial review.

Status of Previous Suggestions

1. ✅ Replace assertions with explicit exceptions — Fixed. The new EnvGenAgent uses explicit raise ValueError instead of assert for the API key check.

2. ✅ Improve JSON extraction error handling — Fixed. The new AgentResponseParseError(ValueError) custom exception is exactly what I suggested, including raw response preview (capped at 500 chars for readability).

3. ℹ️ Retry logic — Not implemented, but this was a nice-to-have suggestion. Callers can wrap as needed.

4. ✅ Add test coverage for LLMAgent/EnvGenAgent — Fixed. The new test_env_gen_agent.py (399 lines) comprehensively covers:

   • __init__ argument/env-var precedence and missing-key guard
   • _extract_json edge cases (fenced, prose, malformed, unbalanced braces, truncated error messages)
   • generate_spec / ping with mocked OpenAI client
   • System prompt coverage for RelationKind/TaskKind literals

5. ⚪ Resolver silent failure documentation — N/A. The entire resolver module was removed in this refactor. The PR now focuses on the agent layer; resolution is presumably handled elsewhere.

New Changes Observed

• Module renamed: llm_env_gen → agentic_env_gen
• Schema renamed: SceneSpec → EnvIntentSpec (with new reasoning field for chain-of-thought)
• Added CI job test_agent_remote_e2e with proper secret validation
• Added pydantic>=2.0 as hard dependency in setup.py
• Removed to_dict() / to_yaml() from ArenaEnvGraphSpec — validation now runs only in from_dict()
• Added at_pose to RelationKind literal

No new issues found in the incremental changes. The refactor improves the codebase significantly.

---

Update (commit fb42596): Major refactor to OpenAI-compatible structured outputs.

Changes Observed

1. Structured outputs migration: The agent now uses response_format={"type": "json_schema", ...} with strict mode, eliminating the need for prose-parsing fallback. This is a significant reliability improvement — the wire guarantees a valid JSON envelope.

2. New structured_output_utils.py module: Well-designed extraction of reusable utilities:

   • build_strict_schema() / apply_strict_constraints() — munges pydantic schemas for OpenAI/Bedrock strict mode (adds additionalProperties: false, populates required, strips default keys)
   • extract_response_text() — handles NVIDIA DeepSeek's reasoning_content channel quirk
   • check_structured_output_support() — deployment validator probe with rich diagnostics

3. Removed code: _extract_json() and AgentResponseParseError removed — no longer needed since structured outputs guarantee the JSON envelope. Clean deletion.

4. Constructor-time fail-fast: Two-stage validation on EnvGenAgent.__init__:

   • First: cheap ping() to verify connectivity/auth
   • Then: full check_structured_output_support() probe with the real schema

   This surfaces capability failures at construction time rather than mid-pipeline.

5. Test coverage: Comprehensive — test_structured_output_utils.py (421 lines) covers:

   • Schema munging for strict mode (``, nested objects, idempotency)
   • Channel extraction fallback (content → reasoning_content → empty)
   • Probe failure modes (4xx API error, empty envelope, parse error, validation error)
   • Live endpoint test gated by @pytest.mark.agent_remote_e2e

Notes

• The json.loads(text, strict=False) call in generate_spec is a good workaround for DeepSeek-v4-flash's unescaped control characters — documented via inline comment.
• Good defensive caching of self._spec_schema to avoid re-walking the schema on every call.
• The system prompt was correctly trimmed to remove the now-redundant "emit ONLY JSON" instruction.

No new issues found. This is a clean, well-tested migration to structured outputs.

---

Update (commit f92fa73): Two defensive improvements.

1. Allow empty tasks list in EnvIntentSpec: Removed the @model_validator that enforced non-empty tasks. Empty tasks now legitimately maps to NoTask for static sandbox environments. The field description clearly documents this behavior.

2. Handle empty choices array in structured output validation: Added defensive check in check_structured_output_support() for when providers (Azure content-filter, Bedrock guardrails) return HTTP 200 with an empty choices list — preventing IndexError and surfacing it cleanly as a parse_error with response_route="empty".

Both changes include test updates. No issues found.

[greptile-apps]

Greptile Summary

This PR introduces the first stage of the agentic environment-generation pipeline: an EnvGenAgent class that calls an OpenAI-compatible structured-outputs endpoint (NVIDIA hosted by default) and parses the response into a validated EnvIntentSpec Pydantic model.

• EnvGenAgent — constructor runs two fail-fast wire checks (cheap ping then a structured-output probe with the full EnvIntentSpec schema) so misconfigured models are rejected before generate_spec is ever called; generate_spec uses response_format=json_schema exclusively, eliminating the prose-extraction fallback from the predecessor module.
• structured_output_utils — schema munging, the NVIDIA DeepSeek reasoning_content channel fallback, and a deployment-validator probe that correctly guards against empty choices lists — though the sibling ping() function has the same gap and is not yet guarded.
• CI & deps — a new agent_remote_e2e job gates live tests behind a secret-presence check; openai and pydantic>=2.0 are added to RUNTIME_DEPS.

Confidence Score: 4/5

Safe to merge with one small fix: ping() needs the same empty-choices guard that was already applied to check_structured_output_support in this PR.

ping() accesses resp.choices[0] unconditionally. The same empty-choices path that check_structured_output_support was explicitly hardened against in this PR can also hit ping, causing an IndexError to propagate through EnvGenAgent.__init__() as an opaque crash rather than a diagnosable wire failure. The rest of the pipeline is well-reasoned and well-tested.

isaaclab_arena/environments/agentic_env_gen/structured_output_utils.py — the ping() function at line 154.

Important Files Changed

Filename	Overview
isaaclab_arena/environments/agentic_env_gen/structured_output_utils.py	New utility module for structured-output probing; check_structured_output_support correctly guards against empty choices, but ping() does not and will raise IndexError on the same empty-choices API paths.
isaaclab_arena/environments/agentic_env_gen/env_intent_spec.py	New Pydantic schema for LLM output; Relation.params: dict (bare untyped dict) generates an object schema without additionalProperties: false, which may fail strict OpenAI-compatible validators.
isaaclab_arena/environments/agentic_env_gen/env_gen_agent.py	New EnvGenAgent with constructor-time ping + structured-output validation and generate_spec using response_format=json_schema; clean refactor from prose-parsing predecessor.
isaaclab_arena/environments/agentic_env_gen/try_schema.py	Demo/debug script; background override now correctly rewrites both rel.subject and rel.target, fixing the previously-noted unary-relation drop.
isaaclab_arena/tests/test_env_gen_agent.py	Comprehensive unit tests for EnvGenAgent covering constructor probes, request shape, channel fallback, error propagation, and an opt-in live E2E test.
isaaclab_arena/tests/test_structured_output_utils.py	Thorough tests for schema munging, channel routing, and all failure modes of check_structured_output_support; the empty-choices test confirms the fix but the parallel gap in ping is untested.
.github/workflows/ci.yml	Adds agent_remote_e2e CI job with explicit secret-presence guard; also excludes the new marker from the default PhysX test run.
setup.py	Adds openai and pydantic>=2.0 to RUNTIME_DEPS, addressing the previously-noted missing pydantic dependency.

Sequence Diagram

sequenceDiagram
    participant U as Caller
    participant A as EnvGenAgent init
    participant P as ping
    participant C as check_support
    participant G as generate_spec
    participant LLM as Inference Endpoint

    U->>A: EnvGenAgent(model, base_url)
    A->>P: ping(client, model)
    P->>LLM: completions.create max_tokens 8
    LLM-->>P: response text
    P-->>A: text
    A->>C: check_structured_output_support(client, model, EnvIntentSpec)
    C->>LLM: completions.create response_format json_schema
    LLM-->>C: EnvIntentSpec JSON sample
    C-->>A: StructuredOutputSupport supported True
    A-->>U: EnvGenAgent instance

    U->>G: generate_spec(prompt, catalog_text)
    G->>LLM: completions.create response_format json_schema
    LLM-->>G: JSON envelope
    G->>G: extract_response_text
    G->>G: json loads strict False
    G->>G: EnvIntentSpec model_validate
    G-->>U: EnvIntentSpec plus raw text

Loading

_{Reviews (9): Last reviewed commit: "Allow empty tasks list in EnvIntentSpec" | Re-trigger Greptile}

[xyao-nv]

How about optional ones, like lightings, hdr image?

[xyao-nv]

Much thorough than the last one. Most comments are to make it more robust and less prone to hallucinations. Feel free to leave some as TODOs.

And one major ask is to remove LLM, and call it Agent. Classes could be called as EnvGenAgent, EnvIntentSpec

[greptile-apps]

ping() accesses resp.choices[0] unconditionally. Some API error paths (e.g. Azure content-filter trips, Bedrock guardrail rejections, and certain rate-limit responses) return HTTP 200 with an empty choices list, making this line raise IndexError — not an openai-client exception. That breaks the documented contract ("Raises: Any exception raised by the underlying openai client") and surfaces as an opaque crash from EnvGenAgent.__init__() rather than a diagnosable ping failure. check_structured_output_support in the same file was explicitly fixed in this PR to guard against this case; ping needs the same treatment.

Suggested change

	resp = client.chat.completions.create(
	model=model,
	messages=[{"role": "user", "content": "Respond with exactly: OK"}],
	temperature=0,
	max_tokens=8,
	)
	return resp.choices[0].message.content or ""
	resp = client.chat.completions.create(
	model=model,
	messages=[{"role": "user", "content": "Respond with exactly: OK"}],
	temperature=0,
	max_tokens=8,
	)
	choices = getattr(resp, "choices", None) or []
	if not choices:
	return ""
	return choices[0].message.content or ""