-
Notifications
You must be signed in to change notification settings - Fork 285
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
[Feature] Categorical encoding for action space #593
Conversation
Codecov Report
@@ Coverage Diff @@
## main #593 +/- ##
==========================================
+ Coverage 87.24% 87.39% +0.14%
==========================================
Files 122 124 +2
Lines 22532 22784 +252
==========================================
+ Hits 19658 19911 +253
+ Misses 2874 2873 -1
Flags with carried forward coverage won't be shown. Click here to find out more.
📣 We’re building smart automated test selection to slash your CI/CD build times. Learn more |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
We'd like to add the new classes to the docs (this is not done automatically atm).
Just add them in docs/source/data.
torchrl/_utils.py
Outdated
key (str): name of the environment variable. | ||
""" | ||
val = os.environ.get(key, False) | ||
if val in ("0", "False", False): |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
if we consider the option False
why not plain 0
?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I changed it to "False"
by default. Environment variables cannot be set to something other than str
.
torchrl/data/tensor_specs.py
Outdated
def __init__( | ||
self, | ||
n: int, | ||
shape: Optional[torch.Size] = torch.Size((1,)), |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
IIRC we once had a linting issue with this (that seems to have disappeared now).
This is why we usually have a None
and if None
then shape is replaced by Size([1])
.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Noted, changed to that
torchrl/objectives/dqn.py
Outdated
|
||
if _CATEGORICAL_ACTION_ENCODING: | ||
batch_size = action.size(0) | ||
pred_val_index = pred_val[range(batch_size), action.squeeze(-1)] |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I guess we'd like to be able to work with batch sizes that are more than unidimensional.
We can use gather
(I think)
x = torch.randn(3, 4, 5)
idx = torch.randint(5, (3, 4))
x.gather(-1, idx)
This will index x
along the last dimension.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
this is not covered by the tests (it should be done in test_costs.py
I believe).
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Reworked indexing here, thx
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Also added tests so should be covered by coverage by now
torchrl/objectives/dqn.py
Outdated
pred_val = td_copy.get("action_value") | ||
pred_val_index = (pred_val * action).sum(-1) | ||
|
||
if _CATEGORICAL_ACTION_ENCODING: |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
This is a bit tricky
What if someone explicitly builds a policy that uses categorical actions, but does not set the env variable?
My view of the env variable is more along the line of "a way of easily change a whole training script from categorical to one hot and vice-versa".
But the reversed dependency should not hold, and the environment variable should only be checked when building the network.
One option could be to infer the type of action by looking at the value_network
: if it has registered the action spec, then we can use it to infer the type of action we will see. Otherwise, we must ask to the user to tell us (via an arg in the constructor) what kind of action we should be expecting.
Open to other suggestions obviously :)
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
What if someone explicitly builds a policy that uses categorical actions, but does not set the env variable?
This is exactly what I meant when mentioned that the current solution couples things a bit too much.
Since actions are passed as pure tensors (even though, wrapped in TensorDict) it is either assumed or configured separately for each class what to do with action
tensors.
Now I can see two options what we can do here:
- Whenever we deal with actions, we either pass
action_space
in the constructor or derive it from some other entity (e.g.value_network
as you proposed). However here we rely on careful configuration of all modules together (in theory, somebody could explicitly configurevalue_network
incompatible with an environment). Another drawback is that we "spreading out" the logic for each case across the code base which can make code a harder to understand (I think in case of one-hot/categorical it is fine, but for future approaches it can blow up). - We can introduce something like
ActionSpec
which would be an extended version of existing specs, though purely for action processing (deriving actions fromvalue_network
, etc.). We instantiate theActionSpec
once and pass its instance everywhere we need encapsulating work with actions. Although, it is hard for me to estimate what interfaceActionSpec
needs in order to cover all possible cases (I mean globally for all RL cases). And this also would require larger scale refactoring.
tldr; I think 1st option seems better choice for this situation and somewhat follows current code base practices (e.g. QValueHook
)
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
So, I reworked it so it can be specified through config and this eliminates the whole problem with working with environment variables. Now we just specify binary flag that will set the mode of handling discrete gym environments and which also be passed as action space to value network and al the hooks.
torchrl/objectives/dqn.py
Outdated
log_ps_a = log_ps_a.view(batch_size, atoms) # log p(s_t, a_t; θonline) | ||
|
||
if _CATEGORICAL_ACTION_ENCODING: | ||
log_ps_a = action_log_softmax[range(batch_size), :, action.squeeze(-1)] |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
same comments as above regarding indexing, usage of env variable and test coverage
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Reworked this as well with a few tweaks since we work with atoms here.
|
||
if _CATEGORICAL_ACTION_ENCODING: | ||
actor_kwargs.update({"action_space": "categorical"}) | ||
out_features = env_specs["action_spec"].space.n |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
perhaps we could have an arg in the function that indicates if categorical has to be used.
If default (e.g. None
) then it falls back on _CATEGORICAL_ACTION_ENCODING
.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I changed it here to be inferred from action spec that is set during the env setup.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Great! Amazing work, if fits in the library very well now. I think we're there.
A couple of things before we merge this
- Can you have a look at the few comments I left?
- Can you have a look at the coverage report (some lines are not covered -- don't worry about things that are not covered but are in the
test
directory).
Oh and one more thing: Can you add the new classes to the doc? Have a look in docs/source/reference
test/test_actors.py
Outdated
|
||
|
||
def test_qvalue_hook_wrong_action_space(): | ||
with pytest.raises(ValueError): |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
let's check that the message match, to make sure we're not capturing the wrong error
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Added a message check (had to make it short since it relies on order of the items in dict)
test/test_actors.py
Outdated
|
||
|
||
def test_distributional_qvalue_hook_wrong_action_space(): | ||
with pytest.raises(ValueError): |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
same as above
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Fixed as in test_qvalue_hook_wrong_action_space
test/test_actors.py
Outdated
) | ||
|
||
|
||
def test_qvalue_hook_wrong_action_space(): |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
could we put all those test_qvalue
under a TestQValue
class?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I guess that if we add a test_actors.py
we should also move some tests there in a future PR
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
could we put all those
test_qvalue
under aTestQValue
class?
Fixed, thanks
@@ -0,0 +1,62 @@ | |||
import os |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Same note to myself: we should move some tests here (e.g. timeit etc)
Awesome! Thank you for the review! |
Description
Added alternative to one-hot encoding for action spaces with categorical features.
Motivation and Context
Implementing feature #538.
Types of changes
What types of changes does your code introduce? Remove all that do not apply:
Checklist
Checklist from #538:
DiscreteTensorSpec
similar toMultOneHotDiscreteTensorSpec
+ testsDistributionalQValueHook
Current implementation seems to couple different parts of TorchRL quite tightly as there is quite a few places where it is assumed that the action space is one-hot. Perhaps, a more generic action-space object that encapsulates manipulation with values and that is shared across losses, actors, hooks, etc. (as a
nn.Module
maybe) might be a worth investment in future.The mode is turned on by
export CATEGORICAL_ACTION_ENCODING=True
command, and it is turned off by default. Another alternative could be to provide it through the config-file, however this, again, would require a bit more generic solution to configure module that will be passed to all the necessary modules.