[Feature] Add info dict key-spec pairs to observation_spec #504
Conversation
Specs are optionally provided via the info_spec property of the info_dict_reader (all readers should inherit from BaseInfoDictReader). When the info_dict_reader is set, GymLikeEnv recovers the specs and adds them to observation_spec.
| spec = {} | ||
|
|
||
| self._info_spec = { | ||
| key: spec.get(key, UnboundedContinuousTensorSpec()) for key in self.keys |
There was a problem hiding this comment.
MyPy gives me a warning here because TensorSpec has abstract method index which UnboundedContinuousTensorSpec does not define.
vmoens
left a comment
There was a problem hiding this comment.
Wonderful,
do you think we could sketch a couple of tests for this?
You could use HalfCheetah-v4 (https://www.gymlibrary.dev/environments/mujoco/half_cheetah/) which comes with a "x_position" in the info dict
| assert "x_position" in env.observation_spec.keys() | ||
| assert isinstance(env.observation_spec["x_position"], UnboundedContinuousTensorSpec) |
There was a problem hiding this comment.
is there something else I could be checking here?
There was a problem hiding this comment.
you could do assert env.observation_spec["x_position"].is_in(tensordict["x_position"])
There was a problem hiding this comment.
it's essentially testing that x_position is what it's supposed to be
Another thing we could do is set up a screwed up tensor spec (e.g. one that has wrong shape or wrong dtype) and check that this assertion raises an error
There was a problem hiding this comment.
I just added some additional tests as you suggested, although it wasn't raising an error at any stage, just returning false for
env.observation_spec["x_position"].is_in(tensordict["x_position"])Would you have expected an exception at some point?
There was a problem hiding this comment.
no but yes if you put an assert if front of it :p
0842a6b to
f8e73d7
Compare
vmoens
left a comment
There was a problem hiding this comment.
LGTM, thanks for this contribution!
Description
Specs are optionally provided via the
info_specproperty of theinfo_dict_reader(all readers should inherit fromBaseInfoDictReader). When theinfo_dict_readeris set,GymLikeEnvrecovers the specs and adds them toobservation_spec.I would appreciate some guidance on how / where best to test these changes.
Motivation and Context
Closes #453
Types of changes
What types of changes does your code introduce? Remove all that do not apply:
Checklist
Go over all the following points, and put an
xin all the boxes that apply.If you are unsure about any of these, don't hesitate to ask. We are here to help!