unilab.algos.torch.common.ane_wrapper.ANEActorWrapper

class unilab.algos.torch.common.ane_wrapper.ANEActorWrapper[source]

Bases: object

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

Methods

__init__(actor_model, obs_dim, action_dim)

predict(obs_np)

Run inference on ANE.

__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

__call__(obs_np)[source]

Call self as a function.