Question About If PPO Training Will Work

Rohan_Jay: <p>Hi. I am relatively inexperienced in PyTorch so I apologize if this question seems silly.<br>
I am working with some friends on implementing PPO from scratch using PyTorch, and part of what I have to do is training the actor and critic network. From what I understand, the cost function of the actor should be the output of the critic. I have attached my current code below.<br>
I have read that PyTorch builds a computational tree of all the computations done to/with a tensor, and <code>.backwards()</code> uses this tree to calculate the gradients and stores it in the tensor. Then <code>.step()</code>, for the tensor passed in(in this case <code>actor.parameters</code>) uses the stored gradients to update the tensor.<br>
However, since I am training both the actor and critic, and am using the critic to calculate the actor’s loss, would this cause a conflict? Would the critic contain the gradients for the actor’s training as well, interfering with the critic’s training?</p>
<pre><code class="lang-auto">def train(actor, critic, env):
	criterion = torch.nn.MSELoss(reduction = 'sum')
	actor_optimizer = torch.optim.SGD(actor.parameters())
	critic_optimizer = torch.optim.SGD(critic.parameters())
	for batch in range(1000):
		env.reset()
		the_action = env.action_space.sample()
	        for t in range(1000):
		          state, act, reward = get_state(env, the_action)
		          critic_pred = critic(the_action)
                         critic_loss = criterion(reward, critic_pred)
                         actor_loss = critic(the_action)#Here

                         #Here
                         actor_optimizer.zero_grad()
                         actor_loss.backward()
                         actor_optimizer.step()

                         critic_optimizer.zero_grad()
                         critic_loss.backward()
                         critic_optimizer.step()

                         the_action = actor(state)
</code></pre>

KFrank: <p>Hi Oliver (and Rohan)!</p>
<p>Without commenting on the entirety of Rohan’s code, your analysis is not correct.</p>
<aside class="quote no-group" data-username="Oliver197Olsen" data-post="2" data-topic="221899" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/o/13edae/48.png" class="avatar"> Oliver197Olsen:</div>
<blockquote>
<p>Your current actor_loss = critic(the_action) line will cause gradients from the actor’s backward() call to accumulate on the critic’s parameters</p>
</blockquote>
</aside>
<p>This part is true – gradients will be accumulated into <code>critic</code>.</p>
<aside class="quote no-group">
<blockquote>
<p>contaminating them before the critic’s own backward() call.</p>
</blockquote>
</aside>
<p>However, this is incorrect because Rohan calls <code>critic_optimizer.zero_grad()</code><br>
before calling <code>critic_loss.backward()</code>.  Any “contamination” from the call to<br>
<code>actor_loss.backward()</code> gets zeroed out before <code>critic</code>’s gradients get correctly<br>
populated by the call to <code>critic_loss.backward()</code>.</p>
<p>One last comment:</p>
<p>Rohan creates two copies of <code>critic</code>’s computation graph – once<br>
with <code>critic_pred = critic(the_action)</code> and then again with<br>
<code>actor_loss = critic(the_action)</code>.</p>
<p>This matters because (as written) the call to <code>actor_loss.backward()</code> releases the<br>
<code>actor_loss</code> copy of <code>critic</code>’s computation graph.  But that’s okay, because the call<br>
to <code>critic_loss.backward()</code> uses the <code>critic_pred</code> copy of the computation graph<br>
and it hasn’t been released yet.</p>
<p>So, as it stands, I think (this part of)  Rohan’s code should work.</p>
<p>Best.</p>
<p>K. Frank</p>