(self, num_envs, num_transitions_per_env, device='cuda')
| 22 | class RolloutStorage(nn.Module): |
| 23 | |
| 24 | def __init__(self, num_envs, num_transitions_per_env, device='cuda'): |
| 25 | |
| 26 | super().__init__() |
| 27 | |
| 28 | self.device = device |
| 29 | |
| 30 | self.num_transitions_per_env = num_transitions_per_env |
| 31 | self.num_envs = num_envs |
| 32 | |
| 33 | # rnn |
| 34 | # self.saved_hidden_states_a = None |
| 35 | # self.saved_hidden_states_c = None |
| 36 | |
| 37 | self.step = 0 |
| 38 | self.stored_keys = list() |
| 39 | |
| 40 | def register_key(self, key: str, shape=(), dtype=torch.float): |
| 41 | # This class was partially copied from https://github.com/NVlabs/ProtoMotions/blob/94059259ba2b596bf908828cc04e8fc6ff901114/phys_anim/agents/utils/data_utils.py |
nothing calls this directly
no outgoing calls
no test coverage detected