Adding Truncated Quantile Critics

[Palamabron]

Hi everyone! I am modifying tmrl (my repo) by adding custom interface, new model, new inputs from my custom plugin - link to plugin here and more wandb features like visualizing model architecture and gradients. I am trying to add TQC (Truncated Quantile Critics) (paper here). However, as in TQC Q-Functions output is bigger than I get dimensionality problem that I am not able to fix:
(...) backup = r + (1 - d) * self.gamma * (sorted_z_part - alpha * next_log_pi) RuntimeError: The size of tensor a (46) must match the size of tensor b (128) at non-singleton dimension 1
Add this the fragment of code of 'custom_algorithms.py' where the problem occurs (tensor dimensions in the comments):

# in this case batch size=128 and 2*[quantiles - top_quantiles_to_drop] = 46
with torch.no_grad():
            new_next_action, next_log_pi = self.model.actor(o2)  # Tensor(batch x 3), Tensor(batch)
            # Compute and cut quantiles at the next state
            next_q1 = self.model_target.q1(o2, new_next_action)  # Tensor(batch x quantiles)
            next_q2 = self.model_target.q2(o2, new_next_action)  # Tensor(batch x quantiles)
            next_z = torch.stack((next_q1, next_q2), dim=1)  # Tensor(batch x nets x quantiles)
            sorted_z, _ = torch.sort(next_z.reshape(batch_size, -1))
            sorted_z_part = sorted_z[:, :self.quantiles_total - self.top_quantiles_to_drop] # Tensor(batch x 2*[quantiles - top_quantiles_to_drop])
            # alpha: Tensor (1,)
            backup = r + (1 - d) * self.gamma * (sorted_z_part - alpha * next_log_pi)

Code of PyTorch TQC is available here. Could someone help me fix this issue or give advice what should I do?

[yannbouteiller]

Hi, this sounds like a cool project!

I am not very familiar with the TQC algorithm but it looks promising, it would be cool if you could push the implementation to tmrl when this works :)

I see that SB3-contrib also has an implementation of TQC

For the error you get, it looks like one of your tensors (r, d, sorted_z_part, next_log_pi) is missing the batch dimension for some reason?

[Palamabron]

Thanks for a different source! I will definitely check it out. With TQC from SamsungLabs I had problem to solve dependencies, so I didn't manage to make it work :(
When I finish (in the next few months) I will create a Readme to point out every major change that I made and some pull requests. :) Furthermore, I am going to publish a paper about TQC and my model for further explanations of how everything works.
In 144 hours I will check if any of my tensors is missing a batch dimension and I will debug StableBaselines3 version of TQC to better understand how it works. I will notify you as soon as I will know anything more about my issue.

[Palamabron]

Hi, this really helped me a lot to find out why it is not working. Thanks for your help. In coming months I will make a pull request with all my changes of the code :)