Trait border_tch_agent::model::ModelBase

pub trait ModelBase {
    fn backward_step(&mut self, loss: &Tensor);
    fn get_var_store_mut(&mut self) -> &mut VarStore;
    fn get_var_store(&self) -> &VarStore;
    fn save<T: AsRef<Path>>(&self, path: T) -> Result<()>;
    fn load<T: AsRef<Path>>(&mut self, path: T) -> Result<()>;
}

Base interface.

Required methods

fn backward_step(&mut self, loss: &Tensor)

Trains the network given a loss.

fn get_var_store_mut(&mut self) -> &mut VarStore

Returns var_store as mutable reference.

fn get_var_store(&self) -> &VarStore

Returns var_store.

fn save<T: AsRef<Path>>(&self, path: T) -> Result<()>

Save parameters of the neural network.

fn load<T: AsRef<Path>>(&mut self, path: T) -> Result<()>

Load parameters of the neural network.

Implementors

impl<F, M> ModelBase for IqnModel<F, M> where
    F: SubModel<Output = Tensor>,
    M: SubModel<Input = Tensor, Output = Tensor>,
    F::Config: DeserializeOwned + Serialize,
    M::Config: DeserializeOwned + Serialize, 

impl<P> ModelBase for Actor<P> where
    P: SubModel<Output = (Tensor, Tensor)>,
    P::Config: DeserializeOwned + Serialize + OutDim, 

impl<Q> ModelBase for DQNModel<Q> where
    Q: SubModel<Output = Tensor>,
    Q::Config: DeserializeOwned + Serialize + OutDim, 

impl<Q> ModelBase for Critic<Q> where
    Q: SubModel2<Output = Tensor>,
    Q::Config: DeserializeOwned + Serialize,