How can I disable gradient updates for some modules in autograd backpropagation?
I'm building a multi-model neural network for reinforcement learning to include an action network, a world model network, and a critic. The idea is train the world model to emulate whatever simulation you are trying to master based on input from the action network and the previous state, to train the critic to maximize the Bellman equation (total reinforcement over time) based on the world model output, and then backpropagate the critic value through the world model to provide gradient targets for training the actions. So - from some state, the action network outputs an action which is fed into the model to generate the next state, and that state feeds into the critic network for evaluation against some goal state.

For all this to work, I must use 3 separate loss functions, one for each network, and they all add something to the gradients in one or more networks but they can be in conflict. For example - to train the world model I use a target from an environmental simulation and for the critic I use a target of the current state reward + discount * next state forecast value. However, to train the a actor I just use the negative critic value as a loss and backpropagate all the way through all three models to calibrate the best action. 

I can make this work without any batching by zeroing out gradients incrementally, but that is inefficient and doesn't let me accumulate gradients for any kind of "time-series batching" optimizer update step. Each model has it's own trainable parameters, but the execution graph flows through all three networks. So inside the calibration loop after firing the networks in sequence:
...

        if self.actor.calibrating:
            self.actor.optimizer.zero_grad()
            #Pick loss For maximizing the value of all actions
            loss = -self.critic.value
            #Backpropagate through all three networks to train actor output
            #How do I stop the critic and model networks from incrementing their gradient values?
            loss.backward(retain_graph=True)
            self.actor.optimizer.step()
        if self.model.calibrating:
            self.model.optimizer.zero_grad()
            #Reduce loss for ambiguous actions
            loss = self.model.get_loss() * self.actor.get_confidence()**2
            #How can I block this from backpropagating through action network?
            loss.backward(retain_graph=True)
            self.model.optimizer.step()
        if self.critic.calibrating:
            self.critic.optimizer.zero_grad()
            #Reduce loss for ambiguous actions
            loss = self.critic.get_loss(self.goal) * self.actor.get_confidence()**2
            #How do I stop this from backpropagating through the model and action networks?
            loss.backward(retain_graph=True)
            self.critic.optimizer.step()


...

Finally - my question is in two parts:


How can I temporarily stop loss.backward() at a given layer without detaching it forever?
How can I block loss.backward() from updating some gradients where I'm just flowing through a model to get gradients for another model?


Got this figured out thanks to a suggestion from a colleague to try the requires_grad setting. (I had assumed that would break the execution graph, but it doesn't)

So - to answer my own two questions:


If you calibrate the chained models in the correct order, you can detach them one at a time so that loss.backward() doesn't run over models that aren't needed. I was thinking that this would break the graph but... this is Pytorch, not Tensorflow 1.x and the graph is regenerated on every forward pass anyway. Silly me for missing this yesterday.
If you set requires_grad to False for a model (or a layer or an individual weight) then loss.backward() will STILL traverse the entire connected graph but it will leave those individual gradients as they were while still setting any gradients earlier in the graph. Exactly what I wanted.


This code works to minimize the execution of unnecessary graph traversals and gradient updates. I still need to refactor it for staggered updates over time so that it can accumulate gradients for several cycles before stepping the optimizers, but this definitely works as intended.

#Step through all models in a chain to create gradient paths from critic back through the world model, to the actor.
    def step(self):
        #Get the current state from the simulation
        state = self.world.state
        #Fire the actor to select a softmax action.
        self.actor(state)
        #run the world simulation on that action.
        self.world.step(self.actor.action)
        #Combine the action and starting state as input to the world model.
        if self.actor.calibrating:
            action_state = torch.cat([self.actor.value, state], dim=0)
        else:
            #Push softmax action closer to 1.0
            action_state = torch.cat([self.actor.hard_value, state], dim=0)
        #Run the model and then the critic on the action_state
        self.critic(self.model(action_state))
        if self.actor.calibrating:
            self.actor.optimizer.zero_grad()
            self.model.requires_grad = False
            self.critic.requires_grad = False
            #Pick loss For maximizing the value of the action choice
            loss = -self.critic.value * self.actor.get_confidence()
            loss.backward(retain_graph=True)
            self.actor.optimizer.step()
        if self.model.calibrating:
            #Don't need to backpropagate through actor again
            self.actor.value.detach_()
            self.model.optimizer.zero_grad()
            self.model.requires_grad = True
            #Reduce loss for ambiguous actions
            loss = self.model.get_loss() * self.actor.get_confidence()**2
            loss.backward(retain_graph=True)
            self.model.optimizer.step()
        if self.critic.calibrating:
            #Don't need to backpropagate through the model or actor again
            self.model.value.detach_()
            self.critic.optimizer.zero_grad()
            self.critic.requires_grad = True
            #Reduce loss for ambiguous actions
            loss = self.critic.get_loss(self.goal) * self.actor.get_confidence()**2
            loss.backward(retain_graph=True)
            self.critic.optimizer.step()

