tinygrad/tinygrad
2
0
Fork 3
mirror of https://github.com/tinygrad/tinygrad.git synced 2026-08-29 13:16:09 +00:00
Code Issues Packages Projects Releases Wiki Activity

Changed DEVICE to Device.DEFAULT in deep_determinist_policy_gradient (#1715)

Browse Source 
* added device in optim and deep

* oops forgot to del print code

* use Device.DEFAULT instead

* removed device
This commit is contained in:
geohotstan
2023-08-31 07:08:51 -07:00
committed by GitHub
parent b5cf274da3
commit 94b1257f5e
1 changed files with 8 additions and 14 deletions
Download Patch File Download Diff File Expand all files Collapse all files
examples/deep_deterministic_policy_gradient.py
+8 -14
View File
@@ -10,9 +10,6 @@ import numpy as np
import gym
DEVICE = "GPU" if getenv("GPU") else "CPU"
class Actor:
def __init__(self, num_actions: int, num_states: int, hidden_size: Tuple[int, int] = (400, 300)):
self.l1 = Tensor.glorot_uniform(num_states, hidden_size[0])
@@ -73,11 +70,11 @@ class Buffer:
record_range = min(self.buffer_counter, self.buffer_capacity)
batch_indices = np.random.choice(record_range, self.batch_size)
state_batch = Tensor(self.state_buffer[batch_indices], device=DEVICE, requires_grad=False)
action_batch = Tensor(self.action_buffer[batch_indices], device=DEVICE, requires_grad=False)
reward_batch = Tensor(self.reward_buffer[batch_indices], device=DEVICE, requires_grad=False)
next_state_batch = Tensor(self.next_state_buffer[batch_indices], device=DEVICE, requires_grad=False)
done_batch = Tensor(self.done_buffer[batch_indices], device=DEVICE, requires_grad=False)
state_batch = Tensor(self.state_buffer[batch_indices], requires_grad=False)
action_batch = Tensor(self.action_buffer[batch_indices], requires_grad=False)
reward_batch = Tensor(self.reward_buffer[batch_indices], requires_grad=False)
next_state_batch = Tensor(self.next_state_buffer[batch_indices], requires_grad=False)
done_batch = Tensor(self.done_buffer[batch_indices], requires_grad=False)
return state_batch, action_batch, reward_batch, next_state_batch, done_batch
@@ -92,7 +89,6 @@ class GaussianActionNoise:
np.random.default_rng()
.normal(self.mean, self.std_dev, size=self.mean.shape)
.astype(np.float32),
device=DEVICE,
requires_grad=False,
)
@@ -158,9 +154,6 @@ class DeepDeterministicPolicyGradient:
target_actor_params = get_parameters(self.target_actor)
target_critic_params = get_parameters(self.target_critic)
if DEVICE == "GPU":
[x.gpu_() for x in actor_params + critic_params + target_actor_params + target_critic_params]
self.actor_optimizer = optim.Adam(actor_params, lr_actor)
self.critic_optimizer = optim.Adam(critic_params, lr_critic)
@@ -207,7 +200,7 @@ class DeepDeterministicPolicyGradient:
target_actions = self.target_actor.forward(next_state_batch, self.max_action)
y = reward_batch + self.gamma * self.target_critic.forward(
next_state_batch, target_actions.detach()
) * (Tensor.ones(*done_batch.shape, device=DEVICE, requires_grad=False) - done_batch)
) * (Tensor.ones(*done_batch.shape, requires_grad=False) - done_batch)
self.critic_optimizer.zero_grad()
critic_value = self.critic.forward(state_batch, action_batch)
@@ -236,7 +229,8 @@ if __name__ == "__main__":
done = False
while not done:
prev_state = Tensor(prev_state, device=DEVICE, requires_grad=False)
print(episode)
prev_state = Tensor(prev_state, requires_grad=False)
action = agent.choose_action(prev_state)
state, reward, done, _, info = env.step(action) # for older gym versions there is only one bool, so remove _