We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
Assertion error running deep_deterministic_policy_gradient.py example
- Python 3.9.2
- osx M1

` python3 deep_deterministic_policy_gradient.py`

```
cuda backend not available No module named 'pycuda'
gpu backend not available No module named 'pyopencl'
metal backend not available No module named 'Metal'
Traceback (most recent call last):
  File "tinygrad/examples/deep_deterministic_policy_gradient.py", line 229, in <module>
    agent = DeepDeterministicPolicyGradient(env)
  File "tinygrad/examples/deep_deterministic_policy_gradient.py", line 166, in __init__
    self.update_network_parameters(tau=1.0)
  File "tinygrad/examples/deep_deterministic_policy_gradient.py", line 176, in update_network_parameters
    target_param.assign(param * tau + target_param * (1.0 - tau))
  File "tinygrad/tinygrad/tensor.py", line 83, in assign
    assert not x.requires_grad  # self requires_grad is okay?
AssertionError
```

--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `examples/deep_deterministic_policy_gradient.py`
Content:
```
1 from typing import Optional, Tuple
2 from numpy.typing import NDArray
3 
4 from tinygrad.tensor import Tensor
5 from tinygrad.nn import optim
6 from tinygrad.helpers import getenv
7 
8 import numpy as np
9 import gym
10 
11 
12 DEVICE = "GPU" if getenv("GPU") else "CPU"
13 
14 
15 class Actor:
16   def __init__(self, num_actions: int, num_states: int, hidden_size: Tuple[int, int] = (400, 300)):
17     self.l1 = Tensor.glorot_uniform(num_states, hidden_size[0])
18     self.l2 = Tensor.glorot_uniform(hidden_size[0], hidden_size[1])
19     self.mu = Tensor.glorot_uniform(hidden_size[1], num_actions)
20 
21   def forward(self, state: Tensor, upper_bound: float) -> Tensor:
22     out = state.dot(self.l1).relu()
23     out = out.dot(self.l2).relu()
24     out = out.dot(self.mu).tanh()
25     output = out * upper_bound
26 
27     return output
28 
29 
30 class Critic:
31   def __init__(self, num_inputs: int, hidden_size: Tuple[int, int] = (400, 300)):
32     self.l1 = Tensor.glorot_uniform(num_inputs, hidden_size[0])
33     self.l2 = Tensor.glorot_uniform(hidden_size[0], hidden_size[1])
34     self.q = Tensor.glorot_uniform(hidden_size[1], 1)
35 
36   def forward(self, state: Tensor, action: Tensor) -> Tensor:
37     inputs = state.cat(action, dim=1)
38     out = inputs.dot(self.l1).relu()
39     out = out.dot(self.l2).relu()
40     q = out.dot(self.q)
41 
42     return q
43 
44 
45 class Buffer:
46   def __init__(self, num_actions: int, num_states: int, buffer_capacity: int = 100000, batch_size: int = 64):
47     self.buffer_capacity = buffer_capacity
48     self.batch_size = batch_size
49 
50     self.buffer_counter = 0
51 
52     self.state_buffer = np.zeros((self.buffer_capacity, num_states))
53     self.action_buffer = np.zeros((self.buffer_capacity, num_actions))
54     self.reward_buffer = np.zeros((self.buffer_capacity, 1))
55     self.next_state_buffer = np.zeros((self.buffer_capacity, num_states))
56     self.done_buffer = np.zeros((self.buffer_capacity, 1))
57 
58   def record(
59     self, observations: Tuple[Tensor, NDArray, float, NDArray, bool]
60   ) -> None:
61     index = self.buffer_counter % self.buffer_capacity
62 
63     self.state_buffer[index] = observations[0].detach().numpy()
64     self.action_buffer[index] = observations[1]
65     self.reward_buffer[index] = observations[2]
66     self.next_state_buffer[index] = observations[3]
67     self.done_buffer[index] = observations[4]
68 
69     self.buffer_counter += 1
70 
71   def sample(self) -> Tuple[Tensor, Tensor, Tensor, Tensor, Tensor]:
72     record_range = min(self.buffer_counter, self.buffer_capacity)
73     batch_indices = np.random.choice(record_range, self.batch_size)
74 
75     state_batch = Tensor(self.state_buffer[batch_indices], device=DEVICE, requires_grad=False)
76     action_batch = Tensor(self.action_buffer[batch_indices], device=DEVICE, requires_grad=False)
77     reward_batch = Tensor(self.reward_buffer[batch_indices], device=DEVICE, requires_grad=False)
78     next_state_batch = Tensor(self.next_state_buffer[batch_indices], device=DEVICE, requires_grad=False)
79     done_batch = Tensor(self.done_buffer[batch_indices], device=DEVICE, requires_grad=False)
80 
81     return state_batch, action_batch, reward_batch, next_state_batch, done_batch
82 
83 
84 class GaussianActionNoise:
85   def __init__(self, mean: NDArray, std_deviation: NDArray):
86     self.mean = mean
87     self.std_dev = std_deviation
88 
89   def __call__(self) -> Tensor:
90     return Tensor(
91         np.random.default_rng()
92         .normal(self.mean, self.std_dev, size=self.mean.shape)
93         .astype(np.float32),
94         device=DEVICE,
95         requires_grad=False,
96     )
97 
98 
99 class DeepDeterministicPolicyGradient:
100   """Deep Deterministic Policy Gradient (DDPG).
101 
102   https://arxiv.org/pdf/1509.02971.pdf
103 
104   Args:
105       env: The environment to learn from.
106       lr_actor: The learning rate of the actor.
107       lr_critic: The learning rate of the critic.
108       gamma: The discount factor.
109       buffer_capacity: The size of the replay buffer.
110       tau: The soft update coefficient.
111       hidden_size: The number of neurons in the hidden layers of the actor and critic networks.
112       batch_size: The minibatch size for each gradient update.
113       noise_stddev: The standard deviation of the exploration noise.
114 
115   Note:
116       In contrast to the original paper, actions are already included in the first layer 
117       of the Critic and we use a Gaussian distribution instead of an Ornstein Uhlenbeck 
118       process for exploration noise.
119 
120   """
121 
122   def __init__(
123     self,
124     env: gym.Env,
125     lr_actor: float = 0.001,
126     lr_critic: float = 0.002,
127     gamma: float = 0.99,
128     buffer_capacity: int = 100000,
129     tau: float = 0.005,
130     hidden_size: Tuple[int, int] = (400, 300),
131     batch_size: int = 64,
132     noise_stddev: float = 0.1,
133   ):
134     self.num_states =  env.observation_space.shape[0]
135     self.num_actions =  env.action_space.shape[0]
136     self.max_action =  env.action_space.high.item()
137     self.min_action =  env.action_space.low.item()
138     self.gamma = gamma
139     self.tau = tau
140     self.memory = Buffer(
141         self.num_actions, self.num_states, buffer_capacity, batch_size
142     )
143     self.batch_size = batch_size
144 
145     self.noise = GaussianActionNoise(
146         mean=np.zeros(self.num_actions),
147         std_deviation=noise_stddev * np.ones(self.num_actions),
148     )
149 
150     self.actor = Actor(self.num_actions, self.num_states, hidden_size)
151     self.critic = Critic(self.num_actions + self.num_states, hidden_size)
152     self.target_actor = Actor(self.num_actions, self.num_states, hidden_size)
153     self.target_critic = Critic(self.num_actions + self.num_states, hidden_size)
154 
155     actor_params = optim.get_parameters(self.actor)
156     critic_params = optim.get_parameters(self.critic)
157     target_actor_params = optim.get_parameters(self.target_actor)
158     target_critic_params = optim.get_parameters(self.target_critic)
159 
160     if DEVICE == "GPU":
161       [x.gpu_() for x in actor_params + critic_params + target_actor_params + target_critic_params]
162 
163     self.actor_optimizer = optim.Adam(actor_params, lr_actor)
164     self.critic_optimizer = optim.Adam(critic_params, lr_critic)
165 
166     self.update_network_parameters(tau=1.0)
167 
168   def update_network_parameters(self, tau: Optional[float] = None) -> None:
169     """Updates the parameters of the target networks via 'soft updates'."""
170     if tau is None:
171       tau = self.tau
172 
173     for param, target_param in zip(
174         optim.get_parameters(self.actor), optim.get_parameters(self.target_actor)
175     ):
176       target_param.assign(param * tau + target_param * (1.0 - tau))
177 
178     for param, target_param in zip(
179         optim.get_parameters(self.critic), optim.get_parameters(self.target_critic)
180     ):
181       target_param.assign(param * tau + target_param * (1.0 - tau))
182 
183   def choose_action(self, state: Tensor, evaluate: bool = False) -> NDArray:
184     mu = self.actor.forward(state, self.max_action)
185 
186     if not evaluate:
187       mu = mu.add(self.noise())
188 
189     mu = mu.clip(self.min_action, self.max_action)
190 
191     return mu.detach().numpy()
192 
193   def learn(self) -> None:
194     """Performs a learning step by sampling from replay buffer and updating networks."""
195     if self.memory.buffer_counter < self.batch_size:
196       return
197 
198     (
199         state_batch,
200         action_batch,
201         reward_batch,
202         next_state_batch,
203         done_batch,
204     ) = self.memory.sample()
205  
206     target_actions = self.target_actor.forward(next_state_batch, self.max_action)
207     y = reward_batch + self.gamma * self.target_critic.forward(
208         next_state_batch, target_actions.detach()
209     ) * (Tensor.ones(*done_batch.shape, device=DEVICE, requires_grad=False) - done_batch)
210 
211     self.critic_optimizer.zero_grad()
212     critic_value = self.critic.forward(state_batch, action_batch)
213     critic_loss = y.detach().sub(critic_value).pow(2).mean()
214     critic_loss.backward()
215     self.critic_optimizer.step()
216 
217     self.actor_optimizer.zero_grad()
218     actions = self.actor.forward(state_batch, self.max_action)
219     critic_value = self.critic.forward(state_batch, actions)
220     actor_loss = -critic_value.mean()
221     actor_loss.backward()
222     self.actor_optimizer.step()
223 
224     self.update_network_parameters()
225 
226 
227 if __name__ == "__main__":
228   env = gym.make("Pendulum-v1")
229   agent = DeepDeterministicPolicyGradient(env)
230   num_episodes = 150
231 
232   for episode in range(1, num_episodes+1):
233     cumulative_reward = 0.0
234     prev_state, info = env.reset()  # for older gym versions only state is returned, so remove info
235     done = False
236 
237     while not done:
238       prev_state = Tensor(prev_state, device=DEVICE, requires_grad=False)
239       action = agent.choose_action(prev_state)
240 
241       state, reward, done, _, info = env.step(action)  # for older gym versions there is only one bool, so remove _
242 
243       cumulative_reward += reward
244 
245       agent.memory.record((prev_state, action, reward, state, done))
246       agent.learn()
247 
248       if done:
249         break
250 
251       prev_state = state
252 
253     print(
254         f"Episode {episode}/{num_episodes} - cumulative reward: {cumulative_reward}"
255     )
256
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/examples/deep_deterministic_policy_gradient.py b/examples/deep_deterministic_policy_gradient.py
--- a/examples/deep_deterministic_policy_gradient.py
+++ b/examples/deep_deterministic_policy_gradient.py
@@ -173,12 +173,12 @@
     for param, target_param in zip(
         optim.get_parameters(self.actor), optim.get_parameters(self.target_actor)
     ):
-      target_param.assign(param * tau + target_param * (1.0 - tau))
+      target_param.assign(param.detach() * tau + target_param * (1.0 - tau))
 
     for param, target_param in zip(
         optim.get_parameters(self.critic), optim.get_parameters(self.target_critic)
     ):
-      target_param.assign(param * tau + target_param * (1.0 - tau))
+      target_param.assign(param.detach() * tau + target_param * (1.0 - tau))
 
   def choose_action(self, state: Tensor, evaluate: bool = False) -> NDArray:
     mu = self.actor.forward(state, self.max_action)

{"golden_diff": "diff --git a/examples/deep_deterministic_policy_gradient.py b/examples/deep_deterministic_policy_gradient.py\n--- a/examples/deep_deterministic_policy_gradient.py\n+++ b/examples/deep_deterministic_policy_gradient.py\n@@ -173,12 +173,12 @@\n     for param, target_param in zip(\n         optim.get_parameters(self.actor), optim.get_parameters(self.target_actor)\n     ):\n-      target_param.assign(param * tau + target_param * (1.0 - tau))\n+      target_param.assign(param.detach() * tau + target_param * (1.0 - tau))\n \n     for param, target_param in zip(\n         optim.get_parameters(self.critic), optim.get_parameters(self.target_critic)\n     ):\n-      target_param.assign(param * tau + target_param * (1.0 - tau))\n+      target_param.assign(param.detach() * tau + target_param * (1.0 - tau))\n \n   def choose_action(self, state: Tensor, evaluate: bool = False) -> NDArray:\n     mu = self.actor.forward(state, self.max_action)\n", "issue": "Assertion error running deep_deterministic_policy_gradient.py example\n- Python 3.9.2\r\n- osx M1\r\n\r\n` python3 deep_deterministic_policy_gradient.py`\r\n\r\n```\r\ncuda backend not available No module named 'pycuda'\r\ngpu backend not available No module named 'pyopencl'\r\nmetal backend not available No module named 'Metal'\r\nTraceback (most recent call last):\r\n  File \"tinygrad/examples/deep_deterministic_policy_gradient.py\", line 229, in <module>\r\n    agent = DeepDeterministicPolicyGradient(env)\r\n  File \"tinygrad/examples/deep_deterministic_policy_gradient.py\", line 166, in __init__\r\n    self.update_network_parameters(tau=1.0)\r\n  File \"tinygrad/examples/deep_deterministic_policy_gradient.py\", line 176, in update_network_parameters\r\n    target_param.assign(param * tau + target_param * (1.0 - tau))\r\n  File \"tinygrad/tinygrad/tensor.py\", line 83, in assign\r\n    assert not x.requires_grad  # self requires_grad is okay?\r\nAssertionError\r\n```\n", "before_files": [{"content": "from typing import Optional, Tuple\nfrom numpy.typing import NDArray\n\nfrom tinygrad.tensor import Tensor\nfrom tinygrad.nn import optim\nfrom tinygrad.helpers import getenv\n\nimport numpy as np\nimport gym\n\n\nDEVICE = \"GPU\" if getenv(\"GPU\") else \"CPU\"\n\n\nclass Actor:\n  def __init__(self, num_actions: int, num_states: int, hidden_size: Tuple[int, int] = (400, 300)):\n    self.l1 = Tensor.glorot_uniform(num_states, hidden_size[0])\n    self.l2 = Tensor.glorot_uniform(hidden_size[0], hidden_size[1])\n    self.mu = Tensor.glorot_uniform(hidden_size[1], num_actions)\n\n  def forward(self, state: Tensor, upper_bound: float) -> Tensor:\n    out = state.dot(self.l1).relu()\n    out = out.dot(self.l2).relu()\n    out = out.dot(self.mu).tanh()\n    output = out * upper_bound\n\n    return output\n\n\nclass Critic:\n  def __init__(self, num_inputs: int, hidden_size: Tuple[int, int] = (400, 300)):\n    self.l1 = Tensor.glorot_uniform(num_inputs, hidden_size[0])\n    self.l2 = Tensor.glorot_uniform(hidden_size[0], hidden_size[1])\n    self.q = Tensor.glorot_uniform(hidden_size[1], 1)\n\n  def forward(self, state: Tensor, action: Tensor) -> Tensor:\n    inputs = state.cat(action, dim=1)\n    out = inputs.dot(self.l1).relu()\n    out = out.dot(self.l2).relu()\n    q = out.dot(self.q)\n\n    return q\n\n\nclass Buffer:\n  def __init__(self, num_actions: int, num_states: int, buffer_capacity: int = 100000, batch_size: int = 64):\n    self.buffer_capacity = buffer_capacity\n    self.batch_size = batch_size\n\n    self.buffer_counter = 0\n\n    self.state_buffer = np.zeros((self.buffer_capacity, num_states))\n    self.action_buffer = np.zeros((self.buffer_capacity, num_actions))\n    self.reward_buffer = np.zeros((self.buffer_capacity, 1))\n    self.next_state_buffer = np.zeros((self.buffer_capacity, num_states))\n    self.done_buffer = np.zeros((self.buffer_capacity, 1))\n\n  def record(\n    self, observations: Tuple[Tensor, NDArray, float, NDArray, bool]\n  ) -> None:\n    index = self.buffer_counter % self.buffer_capacity\n\n    self.state_buffer[index] = observations[0].detach().numpy()\n    self.action_buffer[index] = observations[1]\n    self.reward_buffer[index] = observations[2]\n    self.next_state_buffer[index] = observations[3]\n    self.done_buffer[index] = observations[4]\n\n    self.buffer_counter += 1\n\n  def sample(self) -> Tuple[Tensor, Tensor, Tensor, Tensor, Tensor]:\n    record_range = min(self.buffer_counter, self.buffer_capacity)\n    batch_indices = np.random.choice(record_range, self.batch_size)\n\n    state_batch = Tensor(self.state_buffer[batch_indices], device=DEVICE, requires_grad=False)\n    action_batch = Tensor(self.action_buffer[batch_indices], device=DEVICE, requires_grad=False)\n    reward_batch = Tensor(self.reward_buffer[batch_indices], device=DEVICE, requires_grad=False)\n    next_state_batch = Tensor(self.next_state_buffer[batch_indices], device=DEVICE, requires_grad=False)\n    done_batch = Tensor(self.done_buffer[batch_indices], device=DEVICE, requires_grad=False)\n\n    return state_batch, action_batch, reward_batch, next_state_batch, done_batch\n\n\nclass GaussianActionNoise:\n  def __init__(self, mean: NDArray, std_deviation: NDArray):\n    self.mean = mean\n    self.std_dev = std_deviation\n\n  def __call__(self) -> Tensor:\n    return Tensor(\n        np.random.default_rng()\n        .normal(self.mean, self.std_dev, size=self.mean.shape)\n        .astype(np.float32),\n        device=DEVICE,\n        requires_grad=False,\n    )\n\n\nclass DeepDeterministicPolicyGradient:\n  \"\"\"Deep Deterministic Policy Gradient (DDPG).\n\n  https://arxiv.org/pdf/1509.02971.pdf\n\n  Args:\n      env: The environment to learn from.\n      lr_actor: The learning rate of the actor.\n      lr_critic: The learning rate of the critic.\n      gamma: The discount factor.\n      buffer_capacity: The size of the replay buffer.\n      tau: The soft update coefficient.\n      hidden_size: The number of neurons in the hidden layers of the actor and critic networks.\n      batch_size: The minibatch size for each gradient update.\n      noise_stddev: The standard deviation of the exploration noise.\n\n  Note:\n      In contrast to the original paper, actions are already included in the first layer \n      of the Critic and we use a Gaussian distribution instead of an Ornstein Uhlenbeck \n      process for exploration noise.\n\n  \"\"\"\n\n  def __init__(\n    self,\n    env: gym.Env,\n    lr_actor: float = 0.001,\n    lr_critic: float = 0.002,\n    gamma: float = 0.99,\n    buffer_capacity: int = 100000,\n    tau: float = 0.005,\n    hidden_size: Tuple[int, int] = (400, 300),\n    batch_size: int = 64,\n    noise_stddev: float = 0.1,\n  ):\n    self.num_states =  env.observation_space.shape[0]\n    self.num_actions =  env.action_space.shape[0]\n    self.max_action =  env.action_space.high.item()\n    self.min_action =  env.action_space.low.item()\n    self.gamma = gamma\n    self.tau = tau\n    self.memory = Buffer(\n        self.num_actions, self.num_states, buffer_capacity, batch_size\n    )\n    self.batch_size = batch_size\n\n    self.noise = GaussianActionNoise(\n        mean=np.zeros(self.num_actions),\n        std_deviation=noise_stddev * np.ones(self.num_actions),\n    )\n\n    self.actor = Actor(self.num_actions, self.num_states, hidden_size)\n    self.critic = Critic(self.num_actions + self.num_states, hidden_size)\n    self.target_actor = Actor(self.num_actions, self.num_states, hidden_size)\n    self.target_critic = Critic(self.num_actions + self.num_states, hidden_size)\n\n    actor_params = optim.get_parameters(self.actor)\n    critic_params = optim.get_parameters(self.critic)\n    target_actor_params = optim.get_parameters(self.target_actor)\n    target_critic_params = optim.get_parameters(self.target_critic)\n\n    if DEVICE == \"GPU\":\n      [x.gpu_() for x in actor_params + critic_params + target_actor_params + target_critic_params]\n\n    self.actor_optimizer = optim.Adam(actor_params, lr_actor)\n    self.critic_optimizer = optim.Adam(critic_params, lr_critic)\n\n    self.update_network_parameters(tau=1.0)\n\n  def update_network_parameters(self, tau: Optional[float] = None) -> None:\n    \"\"\"Updates the parameters of the target networks via 'soft updates'.\"\"\"\n    if tau is None:\n      tau = self.tau\n\n    for param, target_param in zip(\n        optim.get_parameters(self.actor), optim.get_parameters(self.target_actor)\n    ):\n      target_param.assign(param * tau + target_param * (1.0 - tau))\n\n    for param, target_param in zip(\n        optim.get_parameters(self.critic), optim.get_parameters(self.target_critic)\n    ):\n      target_param.assign(param * tau + target_param * (1.0 - tau))\n\n  def choose_action(self, state: Tensor, evaluate: bool = False) -> NDArray:\n    mu = self.actor.forward(state, self.max_action)\n\n    if not evaluate:\n      mu = mu.add(self.noise())\n\n    mu = mu.clip(self.min_action, self.max_action)\n\n    return mu.detach().numpy()\n\n  def learn(self) -> None:\n    \"\"\"Performs a learning step by sampling from replay buffer and updating networks.\"\"\"\n    if self.memory.buffer_counter < self.batch_size:\n      return\n\n    (\n        state_batch,\n        action_batch,\n        reward_batch,\n        next_state_batch,\n        done_batch,\n    ) = self.memory.sample()\n \n    target_actions = self.target_actor.forward(next_state_batch, self.max_action)\n    y = reward_batch + self.gamma * self.target_critic.forward(\n        next_state_batch, target_actions.detach()\n    ) * (Tensor.ones(*done_batch.shape, device=DEVICE, requires_grad=False) - done_batch)\n\n    self.critic_optimizer.zero_grad()\n    critic_value = self.critic.forward(state_batch, action_batch)\n    critic_loss = y.detach().sub(critic_value).pow(2).mean()\n    critic_loss.backward()\n    self.critic_optimizer.step()\n\n    self.actor_optimizer.zero_grad()\n    actions = self.actor.forward(state_batch, self.max_action)\n    critic_value = self.critic.forward(state_batch, actions)\n    actor_loss = -critic_value.mean()\n    actor_loss.backward()\n    self.actor_optimizer.step()\n\n    self.update_network_parameters()\n\n\nif __name__ == \"__main__\":\n  env = gym.make(\"Pendulum-v1\")\n  agent = DeepDeterministicPolicyGradient(env)\n  num_episodes = 150\n\n  for episode in range(1, num_episodes+1):\n    cumulative_reward = 0.0\n    prev_state, info = env.reset()  # for older gym versions only state is returned, so remove info\n    done = False\n\n    while not done:\n      prev_state = Tensor(prev_state, device=DEVICE, requires_grad=False)\n      action = agent.choose_action(prev_state)\n\n      state, reward, done, _, info = env.step(action)  # for older gym versions there is only one bool, so remove _\n\n      cumulative_reward += reward\n\n      agent.memory.record((prev_state, action, reward, state, done))\n      agent.learn()\n\n      if done:\n        break\n\n      prev_state = state\n\n    print(\n        f\"Episode {episode}/{num_episodes} - cumulative reward: {cumulative_reward}\"\n    )\n", "path": "examples/deep_deterministic_policy_gradient.py"}], "after_files": [{"content": "from typing import Optional, Tuple\nfrom numpy.typing import NDArray\n\nfrom tinygrad.tensor import Tensor\nfrom tinygrad.nn import optim\nfrom tinygrad.helpers import getenv\n\nimport numpy as np\nimport gym\n\n\nDEVICE = \"GPU\" if getenv(\"GPU\") else \"CPU\"\n\n\nclass Actor:\n  def __init__(self, num_actions: int, num_states: int, hidden_size: Tuple[int, int] = (400, 300)):\n    self.l1 = Tensor.glorot_uniform(num_states, hidden_size[0])\n    self.l2 = Tensor.glorot_uniform(hidden_size[0], hidden_size[1])\n    self.mu = Tensor.glorot_uniform(hidden_size[1], num_actions)\n\n  def forward(self, state: Tensor, upper_bound: float) -> Tensor:\n    out = state.dot(self.l1).relu()\n    out = out.dot(self.l2).relu()\n    out = out.dot(self.mu).tanh()\n    output = out * upper_bound\n\n    return output\n\n\nclass Critic:\n  def __init__(self, num_inputs: int, hidden_size: Tuple[int, int] = (400, 300)):\n    self.l1 = Tensor.glorot_uniform(num_inputs, hidden_size[0])\n    self.l2 = Tensor.glorot_uniform(hidden_size[0], hidden_size[1])\n    self.q = Tensor.glorot_uniform(hidden_size[1], 1)\n\n  def forward(self, state: Tensor, action: Tensor) -> Tensor:\n    inputs = state.cat(action, dim=1)\n    out = inputs.dot(self.l1).relu()\n    out = out.dot(self.l2).relu()\n    q = out.dot(self.q)\n\n    return q\n\n\nclass Buffer:\n  def __init__(self, num_actions: int, num_states: int, buffer_capacity: int = 100000, batch_size: int = 64):\n    self.buffer_capacity = buffer_capacity\n    self.batch_size = batch_size\n\n    self.buffer_counter = 0\n\n    self.state_buffer = np.zeros((self.buffer_capacity, num_states))\n    self.action_buffer = np.zeros((self.buffer_capacity, num_actions))\n    self.reward_buffer = np.zeros((self.buffer_capacity, 1))\n    self.next_state_buffer = np.zeros((self.buffer_capacity, num_states))\n    self.done_buffer = np.zeros((self.buffer_capacity, 1))\n\n  def record(\n    self, observations: Tuple[Tensor, NDArray, float, NDArray, bool]\n  ) -> None:\n    index = self.buffer_counter % self.buffer_capacity\n\n    self.state_buffer[index] = observations[0].detach().numpy()\n    self.action_buffer[index] = observations[1]\n    self.reward_buffer[index] = observations[2]\n    self.next_state_buffer[index] = observations[3]\n    self.done_buffer[index] = observations[4]\n\n    self.buffer_counter += 1\n\n  def sample(self) -> Tuple[Tensor, Tensor, Tensor, Tensor, Tensor]:\n    record_range = min(self.buffer_counter, self.buffer_capacity)\n    batch_indices = np.random.choice(record_range, self.batch_size)\n\n    state_batch = Tensor(self.state_buffer[batch_indices], device=DEVICE, requires_grad=False)\n    action_batch = Tensor(self.action_buffer[batch_indices], device=DEVICE, requires_grad=False)\n    reward_batch = Tensor(self.reward_buffer[batch_indices], device=DEVICE, requires_grad=False)\n    next_state_batch = Tensor(self.next_state_buffer[batch_indices], device=DEVICE, requires_grad=False)\n    done_batch = Tensor(self.done_buffer[batch_indices], device=DEVICE, requires_grad=False)\n\n    return state_batch, action_batch, reward_batch, next_state_batch, done_batch\n\n\nclass GaussianActionNoise:\n  def __init__(self, mean: NDArray, std_deviation: NDArray):\n    self.mean = mean\n    self.std_dev = std_deviation\n\n  def __call__(self) -> Tensor:\n    return Tensor(\n        np.random.default_rng()\n        .normal(self.mean, self.std_dev, size=self.mean.shape)\n        .astype(np.float32),\n        device=DEVICE,\n        requires_grad=False,\n    )\n\n\nclass DeepDeterministicPolicyGradient:\n  \"\"\"Deep Deterministic Policy Gradient (DDPG).\n\n  https://arxiv.org/pdf/1509.02971.pdf\n\n  Args:\n      env: The environment to learn from.\n      lr_actor: The learning rate of the actor.\n      lr_critic: The learning rate of the critic.\n      gamma: The discount factor.\n      buffer_capacity: The size of the replay buffer.\n      tau: The soft update coefficient.\n      hidden_size: The number of neurons in the hidden layers of the actor and critic networks.\n      batch_size: The minibatch size for each gradient update.\n      noise_stddev: The standard deviation of the exploration noise.\n\n  Note:\n      In contrast to the original paper, actions are already included in the first layer \n      of the Critic and we use a Gaussian distribution instead of an Ornstein Uhlenbeck \n      process for exploration noise.\n\n  \"\"\"\n\n  def __init__(\n    self,\n    env: gym.Env,\n    lr_actor: float = 0.001,\n    lr_critic: float = 0.002,\n    gamma: float = 0.99,\n    buffer_capacity: int = 100000,\n    tau: float = 0.005,\n    hidden_size: Tuple[int, int] = (400, 300),\n    batch_size: int = 64,\n    noise_stddev: float = 0.1,\n  ):\n    self.num_states =  env.observation_space.shape[0]\n    self.num_actions =  env.action_space.shape[0]\n    self.max_action =  env.action_space.high.item()\n    self.min_action =  env.action_space.low.item()\n    self.gamma = gamma\n    self.tau = tau\n    self.memory = Buffer(\n        self.num_actions, self.num_states, buffer_capacity, batch_size\n    )\n    self.batch_size = batch_size\n\n    self.noise = GaussianActionNoise(\n        mean=np.zeros(self.num_actions),\n        std_deviation=noise_stddev * np.ones(self.num_actions),\n    )\n\n    self.actor = Actor(self.num_actions, self.num_states, hidden_size)\n    self.critic = Critic(self.num_actions + self.num_states, hidden_size)\n    self.target_actor = Actor(self.num_actions, self.num_states, hidden_size)\n    self.target_critic = Critic(self.num_actions + self.num_states, hidden_size)\n\n    actor_params = optim.get_parameters(self.actor)\n    critic_params = optim.get_parameters(self.critic)\n    target_actor_params = optim.get_parameters(self.target_actor)\n    target_critic_params = optim.get_parameters(self.target_critic)\n\n    if DEVICE == \"GPU\":\n      [x.gpu_() for x in actor_params + critic_params + target_actor_params + target_critic_params]\n\n    self.actor_optimizer = optim.Adam(actor_params, lr_actor)\n    self.critic_optimizer = optim.Adam(critic_params, lr_critic)\n\n    self.update_network_parameters(tau=1.0)\n\n  def update_network_parameters(self, tau: Optional[float] = None) -> None:\n    \"\"\"Updates the parameters of the target networks via 'soft updates'.\"\"\"\n    if tau is None:\n      tau = self.tau\n\n    for param, target_param in zip(\n        optim.get_parameters(self.actor), optim.get_parameters(self.target_actor)\n    ):\n      target_param.assign(param.detach() * tau + target_param * (1.0 - tau))\n\n    for param, target_param in zip(\n        optim.get_parameters(self.critic), optim.get_parameters(self.target_critic)\n    ):\n      target_param.assign(param.detach() * tau + target_param * (1.0 - tau))\n\n  def choose_action(self, state: Tensor, evaluate: bool = False) -> NDArray:\n    mu = self.actor.forward(state, self.max_action)\n\n    if not evaluate:\n      mu = mu.add(self.noise())\n\n    mu = mu.clip(self.min_action, self.max_action)\n\n    return mu.detach().numpy()\n\n  def learn(self) -> None:\n    \"\"\"Performs a learning step by sampling from replay buffer and updating networks.\"\"\"\n    if self.memory.buffer_counter < self.batch_size:\n      return\n\n    (\n        state_batch,\n        action_batch,\n        reward_batch,\n        next_state_batch,\n        done_batch,\n    ) = self.memory.sample()\n \n    target_actions = self.target_actor.forward(next_state_batch, self.max_action)\n    y = reward_batch + self.gamma * self.target_critic.forward(\n        next_state_batch, target_actions.detach()\n    ) * (Tensor.ones(*done_batch.shape, device=DEVICE, requires_grad=False) - done_batch)\n\n    self.critic_optimizer.zero_grad()\n    critic_value = self.critic.forward(state_batch, action_batch)\n    critic_loss = y.detach().sub(critic_value).pow(2).mean()\n    critic_loss.backward()\n    self.critic_optimizer.step()\n\n    self.actor_optimizer.zero_grad()\n    actions = self.actor.forward(state_batch, self.max_action)\n    critic_value = self.critic.forward(state_batch, actions)\n    actor_loss = -critic_value.mean()\n    actor_loss.backward()\n    self.actor_optimizer.step()\n\n    self.update_network_parameters()\n\n\nif __name__ == \"__main__\":\n  env = gym.make(\"Pendulum-v1\")\n  agent = DeepDeterministicPolicyGradient(env)\n  num_episodes = 150\n\n  for episode in range(1, num_episodes+1):\n    cumulative_reward = 0.0\n    prev_state, info = env.reset()  # for older gym versions only state is returned, so remove info\n    done = False\n\n    while not done:\n      prev_state = Tensor(prev_state, device=DEVICE, requires_grad=False)\n      action = agent.choose_action(prev_state)\n\n      state, reward, done, _, info = env.step(action)  # for older gym versions there is only one bool, so remove _\n\n      cumulative_reward += reward\n\n      agent.memory.record((prev_state, action, reward, state, done))\n      agent.learn()\n\n      if done:\n        break\n\n      prev_state = state\n\n    print(\n        f\"Episode {episode}/{num_episodes} - cumulative reward: {cumulative_reward}\"\n    )\n", "path": "examples/deep_deterministic_policy_gradient.py"}]}