[Feature] Discrete SAC

[BY571]

Description

Adding a discrete SAC example

Motivation and Context

Current SAC implementation only supports continuous action spaces. This PR will add the option to run a discrete SAC example based on the paper.

Convergence proof tested on CartPole-v1 (wandb)

• I have raised an issue to propose this change (required for new features and bug fixes)

Types of changes

What types of changes does your code introduce? Remove all that do not apply:

• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds core functionality)
• Breaking change (fix or feature that would cause existing functionality to change)
• Documentation (update in the documentation)
• Example (update in the folder of examples)

Checklist

Go over all the following points, and put an x in all the boxes that apply.
If you are unsure about any of these, don't hesitate to ask. We are here to help!

• I have read the CONTRIBUTION guide (required)
• My change requires a change to the documentation.
• I have updated the tests accordingly (required for a bug fix or a new feature).
• I have updated the documentation accordingly.

[vmoens]

LGTM
Before landing:

• can we move the loss to sac.py? I'd rather have them all in the same place if that makes sense?
• Can we add the loss to the doc?
• Is this supposed to work with gSDE? gSDE is not tailored for discrete action spaces AFAICT

[BY571]

LGTM Before landing:

• can we move the loss to sac.py? I'd rather have them all in the same place if that makes sense?
• Can we add the loss to the doc?
• Is this supposed to work with gSDE? gSDE is not tailored for discrete action spaces AFAICT

Do you mean discrete and continuous sac loss in one objective class or having both losses just in the same file?
I'd prefer to have them in the same class, what do you think? Will have a look at it in the coming days.

Will add it to the doc and also take off the gSDE :)

[vmoens]

I'd prefer to have them in the same class
How much control flow would that entail?
Does it save a lot of code?

I was thinking of having them in the same file. If having them in the same class does not create a monster class I'm happy to consider it.

Will it work with v1 and v2?

[BY571]

I was thinking of having them in the same file. If having them in the same class does not create a monster class I'm happy to consider it.

Will it work with v1 and v2?

For now, I just added it to the sac.py file in objectives. As it only works with v2 it might get messy and as you said would probably create a monster class. Let me know what you think.

I also took off the gSDE from the loss and updated the description of the actor_network to be a TensorDictModule.

How can I update the docs?

[vmoens]

LGTM

[vmoens]

LGTM -- let's try to merge this :)

[vmoens]

careful here: the [0] can be the batch size
maybe the last dimension? Or since it's discrete we can check if it's a one hot or a discrete encoding and directly retrieve the number of options from the spec?

[vmoens]

I think it's the only place where we use action_spec
Maybe we could just pass the number of possible actions rather than passing the spec?

[BY571]

Fixed that and adapted the example script due to recent TorchRL changes.
However, now some tests fail but I'm at it and hope to resolve them quickly!

[BY571]

Fixed the example script issues and updated the objective tests as well as they were getting several errors.
Hopefully ready to merge now! :)