(
self,
tensordict,
)
| 310 | |
| 311 | |
| 312 | def _loss_value( |
| 313 | self, |
| 314 | tensordict, |
| 315 | ): |
| 316 | td_copy = tensordict.clone() |
| 317 | |
| 318 | # V(s, a) |
| 319 | with self.value_network_params.to_module(self.value_network): |
| 320 | self.value_network(td_copy) |
| 321 | pred_val = td_copy.get("state_action_value").squeeze(-1) |
| 322 | |
| 323 | # we manually reconstruct the parameters of the actor-critic, where the first |
| 324 | # set of parameters belongs to the actor and the second to the value function. |
| 325 | target_params = TensorDict( |
| 326 | { |
| 327 | "module": { |
| 328 | "0": self.target_actor_network_params, |
| 329 | "1": self.target_value_network_params, |
| 330 | } |
| 331 | }, |
| 332 | batch_size=self.target_actor_network_params.batch_size, |
| 333 | device=self.target_actor_network_params.device, |
| 334 | ) |
| 335 | with target_params.to_module(self.actor_critic): |
| 336 | target_value = self.value_estimator.value_estimate(tensordict).squeeze(-1) |
| 337 | |
| 338 | # Computes the value loss: L2, L1 or smooth L1 depending on `self.loss_function` |
| 339 | loss_value = distance_loss(pred_val, target_value, loss_function=self.loss_function) |
| 340 | td_error = (pred_val - target_value).pow(2) |
| 341 | |
| 342 | return loss_value, td_error, pred_val, target_value |
| 343 | |
| 344 | |
| 345 | ############################################################################### |
nothing calls this directly
no outgoing calls
no test coverage detected