oa::ml::policy

Module published by the oa::ml module.

MdRust public API

Public Types

Public Functions

pub fn evaluate_categorical( logits: &Matrix, action: &Matrix, value: &Matrix, ) -> Result<PolicyResult>

pub fn evaluate_tanh_normal( mean: &Matrix, log_stddev: &Matrix, raw_action: &Matrix, value: &Matrix, minimum: f32, maximum: f32, epsilon: f32, ) -> Result<ContinuousPolicyResult>

pub fn sample_categorical( logits: &Matrix, value: &Matrix, seed: u64, ) -> Result<PolicyResult>

pub fn sample_tanh_normal( mean: &Matrix, log_stddev: &Matrix, value: &Matrix, minimum: f32, maximum: f32, seed: u64, epsilon: f32, ) -> Result<ContinuousPolicyResult>

Function Documentation

pub fn evaluate_categorical(
	logits: &Matrix,
	action: &Matrix,
	value: &Matrix
) -> Result<PolicyResult>

Differentiably evaluate stored categorical actions. Gradients flow through logits, log_probability, and entropy; action indices remain discrete. action and value are zero-copy semantic output aliases, matching OA C++. Errors Returns an error when logits are not FP32 [E,A] with E > 0 and A > 1, actions are not I32 [E], values are not FP32 [E], inputs do not share one engine, or composite lowering fails.

pub fn evaluate_tanh_normal(
	mean: &Matrix,
	log_stddev: &Matrix,
	raw_action: &Matrix,
	value: &Matrix,
	minimum: f32,
	maximum: f32,
	epsilon: f32
) -> Result<ContinuousPolicyResult>

Differentiably evaluate a stored pre-tanh normal action. Log standard deviation is clamped to [-20, 2]. Actions are mapped from [-1, 1] into [minimum, maximum]; log probability includes the donor tanh Jacobian correction, while entropy is the unsquashed diagonal-normal value. raw_action and value are zero-copy semantic output aliases. Errors Returns an error when mean, log-standard-deviation, and raw-action are not matching nonempty FP32 [E,A] matrices, value is not FP32 [E], inputs do not share one engine, bounds/epsilon are invalid, a size exceeds the shader ABI, or composite lowering fails.

pub fn sample_categorical(
	logits: &Matrix,
	value: &Matrix,
	seed: u64
) -> Result<PolicyResult>

Sample and evaluate one categorical action per environment. A nonzero seed is deterministic. Seed zero follows the process-local OA RNG convention. The action is sampled from the same logits used for returned log probabilities and entropy. Errors Returns an error when logits are not FP32 [E,A] with E > 0 and A > 1, values are not FP32 [E] on the same engine, or composite lowering fails.

pub fn sample_tanh_normal(
	mean: &Matrix,
	log_stddev: &Matrix,
	value: &Matrix,
	minimum: f32,
	maximum: f32,
	seed: u64,
	epsilon: f32
) -> Result<ContinuousPolicyResult>

Sample a reparameterized tanh-normal policy. The returned raw action is mean + exp(clamped_log_stddev) * noise, so gradients flow through both distribution parameters while Philox noise stays detached. A nonzero seed is deterministic; seed zero follows OA’s process-local RNG convention. Errors Returns an error when mean/log-standard-deviation are not matching nonempty FP32 [E,A] matrices, value is not FP32 [E] on the same engine, bounds/epsilon are invalid, a size exceeds the shader ABI, or composite lowering fails.