Dynamic session providers for onnxruntime

[Kaweees]

This merge request utilizes onnxruntime.get_available_providers to fetch ONNX Runtime Execution Providers, which will outperform the hard-coded default CPUExecutionProvider in cases where another Execution Provider exsists on the machine from which dimensional is ran. See https://onnxruntime.ai/docs/execution-providers/ for more information

[greptile-apps]

Greptile Overview

Greptile Summary

This PR modernizes ONNX Runtime execution provider selection in the MuJoCo policy module by replacing the hardcoded CPUExecutionProvider with dynamic provider detection using ort.get_available_providers().

Key Changes:

1. Import alias standardization: Changed from rt to ort (line 23), following the widely-adopted community convention
2. Dynamic provider selection: Replaced providers=["CPUExecutionProvider"] with providers=ort.get_available_providers() (line 43), enabling automatic hardware acceleration
3. Observability improvement: Added logging statement (line 44) that displays the loaded policy path and the actual execution providers selected by ONNX Runtime

How It Works:
The get_available_providers() function queries the installed ONNX Runtime package and returns a prioritized list of execution providers. For base installations, this returns ["CPUExecutionProvider"] (identical to the previous hardcoded behavior). For GPU-enabled installations (onnxruntime-gpu), it may return ["CUDAExecutionProvider", "CPUExecutionProvider"] or other accelerated providers like TensorrtExecutionProvider. ONNX Runtime's InferenceSession automatically attempts providers in order and falls back gracefully if initialization fails.

Benefits:

• Performance: Enables automatic CUDA/TensorRT acceleration when available, potentially providing significant speedups for robotics control inference
• Flexibility: Works correctly across different hardware configurations without code changes
• Maintainability: Removes hardcoded assumptions and follows ONNX Runtime best practices
• Transparency: New logging helps diagnose which execution provider is actually being used

Compatibility:
This change is fully backward compatible. On systems without GPU support, the behavior is identical to before (CPU execution). The change aligns with similar patterns in the codebase (e.g., sam_2d_seg.py checks for CUDA availability, and image_embedding.py uses provider lists with fallback logic).

Confidence Score: 5/5

• This PR is safe to merge with high confidence - it's a low-risk improvement following ONNX Runtime best practices
• The change is minimal, well-understood, and follows the recommended ONNX Runtime approach. Using get_available_providers() is the standard pattern that enables automatic hardware acceleration while maintaining backward compatibility. The automatic fallback behavior in InferenceSession ensures that if preferred providers fail, CPU execution is used. The addition of logging improves observability without affecting functionality. No breaking changes or new dependencies are introduced.
• No files require special attention

Important Files Changed

File Analysis

Filename	Score	Overview
dimos/simulation/mujoco/policy.py	5/5	Import alias modernized to 'ort' and provider selection changed to use dynamic get_available_providers() for automatic hardware acceleration detection, with helpful logging added

Sequence Diagram

sequenceDiagram
    participant Model as load_model()
    participant Controller as OnnxController.__init__()
    participant ORT as onnxruntime
    participant Logger as logger
    participant Hardware as Hardware/Drivers

    Model->>Controller: Initialize with policy_path
    Controller->>ORT: get_available_providers()
    ORT->>Hardware: Check available execution providers
    Hardware-->>ORT: Return ["CUDAExecutionProvider", "CPUExecutionProvider"] or ["CPUExecutionProvider"]
    ORT-->>Controller: Available providers list
    
    Controller->>ORT: InferenceSession(policy_path, providers=available_providers)
    ORT->>Hardware: Attempt to initialize providers in order
    alt CUDA available and functional
        Hardware-->>ORT: CUDAExecutionProvider initialized
    else CUDA unavailable
        Hardware-->>ORT: Fallback to CPUExecutionProvider
    end
    ORT-->>Controller: Session created with active provider
    
    Controller->>ORT: get_providers()
    ORT-->>Controller: Actually used provider(s)
    Controller->>Logger: Log policy path and active providers
    Logger-->>Controller: Logged
    Controller-->>Model: Initialized OnnxController

Loading

[greptile-apps]

_{1 file reviewed, 1 comment}

_{Edit Code Review Agent Settings | Greptile}

[Kaweees]

@greptile

[greptile-apps]

_{No files reviewed, no comments}

_{Edit Code Review Agent Settings | Greptile}

[greptile-apps]

_{1 file reviewed, 1 comment}

_{Edit Code Review Agent Settings | Greptile}

[greptile-apps]

Consider adding logging to show which execution provider is actually being used after initialization. This would help with debugging and verifying that GPU acceleration is working when available.

Similar to the pattern used in dimos/agents_deprecated/memory/image_embedding.py:89-91:

self._policy = ort.InferenceSession(policy_path, providers=ort.get_available_providers())
actual_providers = self._policy.get_providers()
logger.info(f"Loaded policy with providers: {actual_providers}")

This is especially useful since get_available_providers() returns all available providers, but the InferenceSession may only successfully initialize with a subset of them.

_{Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!}

[greptile-apps]

_{1 file reviewed, 1 comment}

_{Edit Code Review Agent Settings | Greptile}

[greptile-apps]

The change from hardcoded ["CPUExecutionProvider"] to dynamic ort.get_available_providers() lacks error handling and observability. While the implementation is technically correct (contrary to the previous review thread), there are important considerations:

Issues:

1. No logging: Users won't know which execution provider is actually being used, making debugging and performance analysis difficult
2. No error handling: If provider initialization fails (e.g., CUDA provider available but GPU out of memory), there's no graceful fallback or informative error message
3. Silent performance changes: Different providers have different performance characteristics and numerical precision - users should be informed which one is selected

Recommendation:
Add logging after initialization to show which provider was selected, similar to the pattern used in dimos/agents_deprecated/memory/image_embedding.py:89-91:

self._policy = ort.InferenceSession(policy_path, providers=ort.get_available_providers())
actual_providers = self._policy.get_providers()
# Log: f"Initialized ONNX policy with providers: {actual_providers}"

This would help users understand which execution provider is being used and debug any performance or behavior differences.

_{Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!}

[Kaweees]

@greptile

[greptile-apps]

_{No files reviewed, no comments}

_{Edit Code Review Agent Settings | Greptile}