Onnx

navground.learning.onnx

Exporting

export(policy: BasePolicy, path: PathLike, dynamic_batch_size: bool = True) → None

Export a policy (PyTorch) as onnx.

Parameters:

• policy (BasePolicy) – The policy

• path (PathLike) – Where to save the policy; by convention should have an “.onnx” suffix.

• dynamic_batch_size (bool) – Whether to enable dynamic batch size. If not set, it will export a model that evaluate a single sample.

Return type:

None

Inference

class OnnxPolicy(path: PathLike)

This class implements the navground.learning.types.PolicyPredictor protocol. It loads an onnx model andd then calls onnxruntime.InferenceSession.run() to perform inference.

Parameters:

path (PathLike) – The path

predict(observation: Observation, state: State | None = None, episode_start: EpisodeStart | None = None, deterministic: bool = True) → tuple[Action, State | None]

Perform one inference, calling onnxruntime.InferenceSession.run().

The observation should respect the model specifics, also exposed by observation_space.

Parameters:

• observation (Observation) – The input

• state (State | None) – The state (ignored)

• episode_start (EpisodeStart | None) – The episode start (ignored)

• deterministic (bool) – Whether deterministic (ignored)

Returns:

The output of the inference and a None state

Return type:

tuple[Action, State | None]

property action_space: Box

The action space

property observation_space: Box | Dict

The observation space