Method step

4-atari-hard/3-robustify.py:98–105

One timestep for acting. obs (B,4,105,80), hx (B,gru). Returns logits, value, new hx.

(self, obs, hx, inc_entropy=None)

        return h

    def step(self, obs, hx, inc_entropy=None):
        """One timestep for acting. obs (B,4,105,80), hx (B,gru). Returns
        logits, value, new hx."""
        hx = self.gru(self.features(obs), hx)
        logits = self.pi(hx)
        if inc_entropy is not None:
            logits = torch.where(inc_entropy.unsqueeze(1), logits / 2.0, logits)
        return logits, self.v(hx).squeeze(-1), hx

    def unroll(self, obs_seq, hx0, done_seq):
        """Recompute a (T,B) rollout's logits/values with done-masked GRU state