Trait border_tch_agent::model::ModelBase

pub trait ModelBase {
    // Required methods
    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,