Struct border_tch_agent::iqn::IqnModel

pub struct IqnModel<F, M> where
    F: SubModel<Output = Tensor>,
    M: SubModel<Input = Tensor, Output = Tensor>,
    F::Config: DeserializeOwned + Serialize,
    M::Config: DeserializeOwned + Serialize,  { /* private fields */ }

Constructs IQN output layer, which takes input features and percent points. It returns action-value quantiles.

Implementations

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

pub fn build(
    config: IqnModelConfig<F::Config, M::Config>,
    device: Device
) -> Result<IqnModel<F, M>>

Constructs IqnModel.

pub fn build_with_submodel_configs(
    config: IqnModelConfig<F::Config, M::Config>,
    f_config: F::Config,
    m_config: M::Config,
    device: Device
) -> IqnModel<F, M>

Constructs IqnModel with the given configurations of sub models.

pub fn forward(&self, x: &F::Input, tau: &Tensor) -> Tensor

Returns the tensor of action-value quantiles.

• The shape of psi(x) (feature vector) is [batch_size, feature_dim].
• The shape of tau is [batch_size, n_percent_points].
• The shape of the output is [batch_size, n_percent_points, self.out_dim].

Trait Implementations

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

fn clone(&self) -> Self

Returns a copy of the value.

fn clone_from(&mut self, source: &Self)

Performs copy-assignment from source.

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, 

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

Trains the network given a loss.

fn get_var_store(&self) -> &VarStore

Returns var_store.

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

Returns var_store as mutable reference.

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.