unilab.algos.torch.common.ane_wrapper

ANE (Apple Neural Engine) inference wrapper for Fast SAC.

Classes

ANEActorWrapper

Wrapper to run actor inference on Apple Neural Engine via CoreML.
class unilab.algos.torch.common.ane_wrapper.ANEActorWrapper[source]

Bases: object

Wrapper to run actor inference on Apple Neural Engine via CoreML.

__init__(actor_model, obs_dim, action_dim)[source]
predict(obs_np)[source]

Run inference on ANE.

Parameters:

obs_np – numpy array of shape (batch, obs_dim)

Returns:

numpy array of shape (batch, action_dim)

Return type:

actions