| 91 | |
| 92 | # C++ GPU Class |
| 93 | class VecTaskGPU(VecTask): |
| 94 | def __init__(self, task, rl_device, clip_observations=5.0, clip_actions=1.0): |
| 95 | super().__init__(task, rl_device, clip_observations=clip_observations, clip_actions=clip_actions) |
| 96 | |
| 97 | self.obs_tensor = gymtorch.wrap_tensor(self.task.obs_tensor, counts=(self.task.num_envs, self.task.num_obs)) |
| 98 | self.rewards_tensor = gymtorch.wrap_tensor(self.task.rewards_tensor, counts=(self.task.num_envs,)) |
| 99 | self.resets_tensor = gymtorch.wrap_tensor(self.task.resets_tensor, counts=(self.task.num_envs,)) |
| 100 | |
| 101 | def step(self, actions): |
| 102 | self.task.render(False) |
| 103 | actions_clipped = torch.clamp(actions, -self.clip_actions, self.clip_actions) |
| 104 | actions_tensor = gymtorch.unwrap_tensor(actions_clipped) |
| 105 | |
| 106 | self.task.step(actions_tensor) |
| 107 | |
| 108 | return torch.clamp(self.obs_tensor, -self.clip_obs, self.clip_obs), self.rewards_tensor, self.resets_tensor, [] |
| 109 | |
| 110 | def reset(self): |
| 111 | actions = 0.01 * (1 - 2 * torch.rand([self.task.num_envs, self.task.num_actions], dtype=torch.float32, device=self.rl_device)) |
| 112 | actions_tensor = gymtorch.unwrap_tensor(actions) |
| 113 | |
| 114 | # step the simulator |
| 115 | self.task.step(actions_tensor) |
| 116 | |
| 117 | return torch.clamp(self.obs_tensor, -self.clip_obs, self.clip_obs) |
| 118 | |
| 119 | |
| 120 | # Python CPU/GPU Class |
nothing calls this directly
no outgoing calls
no test coverage detected