Skip to content

RANDAL

RANDAL

Random Network Distillation with Auxiliary Learning (RANDAL).

RANDAL 

RANDAL(policy_model, target_model, input_size, action_size, config: PPOConfig, *, pixel_control: bool = True, intr_coeff: float = 0.5, extr_coeff: float = 1.0, RP: float = 1, VR: float = 1, PC: float = 1, policy_args: dict | None = None, RND_args: dict | None = None, optim: type[Optimizer] = torch.optim.Adam, optim_args: dict | None = None)

Bases: Agent

Source code in rlib/RANDAL/model.py 
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
def __init__(
    self,
    policy_model,
    target_model,
    input_size,
    action_size,
    config: PPOConfig,
    *,
    pixel_control: bool = True,
    intr_coeff: float = 0.5,
    extr_coeff: float = 1.0,
    RP: float = 1,
    VR: float = 1,
    PC: float = 1,
    policy_args: dict | None = None,
    RND_args: dict | None = None,
    optim: type[torch.optim.Optimizer] = torch.optim.Adam,
    optim_args: dict | None = None,
):
    if RND_args is None:
        RND_args = {}
    if policy_args is None:
        policy_args = {}
    super().__init__(config=config)
    self.entropy_coeff = config.entropy_coeff
    self.intr_coeff = intr_coeff
    self.extr_coeff = extr_coeff
    self.pixel_control = pixel_control
    self.action_size = action_size
    self.RP = RP  # reward prediction
    self.VR = VR  # value replay
    self.PC = PC  # pixel control

    self.policy = PPOIntrinsic(
        policy_model,
        input_size,
        action_size,
        config=config,
        extr_coeff=extr_coeff,
        intr_coeff=intr_coeff,
        build_optimiser=False,
        **policy_args,
    )

    device = config.device
    target_size = (
        (1, input_size[1], input_size[2]) if len(input_size) == 3 else input_size
    )  # only use last frame in frame-stack for convolutions

    # randomly weighted and fixed neural network, acts as a random_id for each state
    self.target_model = target_model(target_size, trainable=False, **RND_args).to(device)

    # learns to predict target model
    # i.e. provides rewards based ability to predict a fixed random function, thus behaves as density map of explored areas
    self.predictor_model = target_model(target_size, trainable=True, **RND_args).to(device)

    if pixel_control:
        self.feat_map = torch.nn.Sequential(
            torch.nn.Linear(self.policy.dense_size, 32 * 8 * 8), torch.nn.ReLU()
        ).to(device)
        self.deconv1 = torch.nn.Sequential(
            torch.nn.ConvTranspose2d(32, 32, kernel_size=(3, 3), stride=(1, 1)), torch.nn.ReLU()
        ).to(device)
        self.deconv_advantage = torch.nn.ConvTranspose2d(
            32, action_size, kernel_size=(3, 3), stride=(2, 2)
        ).to(device)
        self.deconv_value = torch.nn.ConvTranspose2d(
            32, 1, kernel_size=(3, 3), stride=(2, 2)
        ).to(device)

    # reward model
    self.r1 = torch.nn.Sequential(
        torch.nn.Linear(self.policy.dense_size, 128), torch.nn.ReLU()
    ).to(device)
    self.r2 = torch.nn.Linear(128, 3).to(device)

    self._build_optimiser(optim=optim, optim_args=optim_args)

pixel_loss 

pixel_loss(Qaux, Qaux_actions, Qaux_target)

Qaux_target temporal difference target for Q_aux

Source code in rlib/RANDAL/model.py 
124
125
126
127
128
129
130
131
132
def pixel_loss(self, Qaux, Qaux_actions, Qaux_target):
    'Qaux_target temporal difference target for Q_aux'
    one_hot_actions = F.one_hot(Qaux_actions.long(), self.action_size)
    pixel_action = one_hot_actions.view([-1, self.action_size, 1, 1])
    Q_aux_action = torch.sum(Qaux * pixel_action, dim=1)
    pixel_loss = 0.5 * torch.mean(
        torch.square(Qaux_target - Q_aux_action)
    )  # l2 loss for Q_aux over all pixels and batch
    return pixel_loss

predictor_loss 

predictor_loss(next_states, state_mean, state_std)

loss for predictor network

Source code in rlib/RANDAL/model.py 
188
189
190
191
192
193
194
def predictor_loss(self, next_states, state_mean, state_std):
    'loss for predictor network'
    next_states, state_mean, state_std = totorch_many(
        next_states, state_mean, state_std, device=self.device
    )
    predictor_loss = self._intr_reward(next_states, state_mean, state_std).mean()
    return predictor_loss

RANDALTrainer 

RANDALTrainer(envs, agent: RANDAL, val_envs, config: RANDALTrainerConfig)

Bases: SyncMultiEnvTrainer

Trainer for the RANDAL agent (RND + UNREAL auxiliary tasks).

Source code in rlib/RANDAL/trainer.py 
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
def __init__(
    self,
    envs,
    agent: RANDAL,
    val_envs,
    config: RANDALTrainerConfig,
):
    super().__init__(envs, agent, val_envs, config=config)

    self.gamma_intr = config.gamma_intr
    self.num_epochs = config.num_epochs
    self.num_minibatches = config.num_minibatches
    self.init_obs_steps = config.init_obs_steps
    self.replay_length = config.replay_length
    self.normalise_obs = config.norm_pixel_reward
    self.pred_prob = 1 / (self.num_envs / 32.0)
    self.state_obs = RunningMeanStd()
    self.forward_filter = RewardForwardFilter(config.gamma_intr)
    self.intr_rolling = RunningMeanStd()
    self.replay = deque([], maxlen=config.replay_length)  # replay length per actor

norm_obs 

norm_obs(obs)

normalise pixel intensity changes by recording min and max pixel observations not using per pixel normalisation because expected image is singular greyscale frame

Source code in rlib/RANDAL/trainer.py 
74
75
76
77
78
def norm_obs(self, obs):
    '''normalise pixel intensity changes by recording min and max pixel observations
    not using per pixel normalisation because expected image is singular greyscale frame
    '''
    return (obs - self.state_min) * (1 / (self.state_max - self.state_min))

RANDALTrainerConfig dataclass 

RANDALTrainerConfig(train_mode: TrainMode = TrainMode.NSTEP, returns: Returns = Returns.NSTEP, total_steps: int = 50000000, nsteps: int = 5, gamma: float = 0.99, lambda_: float = 0.95, validate_freq: int = 1000000, num_val_episodes: int = 50, max_val_steps: int = 10000, log_dir: str = 'logs/', model_dir: str = 'models/', save_freq: int = 0, log_scalars: bool = True, update_target_freq: int = 0, render_freq: int = 0, gamma_intr: float = 0.99, init_obs_steps: int = 600, num_epochs: int = 4, num_minibatches: int = 4, replay_length: int = 2000, norm_pixel_reward: bool = True)

Bases: RNDTrainerConfig

Hyperparameters for :class:RANDALTrainer.

Inherits the RND extra fields and adds the UNREAL replay buffer knobs.