border_tch_agent/cnn/
config.rs

1use crate::util::OutDim;
2use serde::{Deserialize, Serialize};
3
4fn default_skip_linear() -> bool {
5    false
6}
7
8#[allow(clippy::upper_case_acronyms)]
9#[derive(Debug, Deserialize, Serialize, PartialEq, Clone)]
10/// Configuration of [`AtariCnn`](super::AtariCnn).
11///
12/// If `skip_linear` is `true`, `out_dim` is not used.
13pub struct AtariCnnConfig {
14    pub n_stack: i64,
15    pub out_dim: i64,
16    #[serde(default = "default_skip_linear")]
17    pub skip_linear: bool,
18}
19
20impl AtariCnnConfig {
21    /// Constructs [`AtariCnnConfig`]
22    pub fn new(n_stack: i64, out_dim: i64) -> Self {
23        Self {
24            n_stack,
25            out_dim,
26            skip_linear: false,
27        }
28    }
29
30    pub fn skip_linear(mut self, skip_linear: bool) -> Self {
31        self.skip_linear = skip_linear;
32        self
33    }
34}
35
36impl OutDim for AtariCnnConfig {
37    /// Gets output dimension.
38    fn get_out_dim(&self) -> i64 {
39        self.out_dim
40    }
41
42    /// Sets output dimension.
43    fn set_out_dim(&mut self, v: i64) {
44        self.out_dim = v;
45    }
46}
border_tch_agent/cnn/config.rs

border_tch_agent/cnn/
config.rs
