Skip to content

Commit

Permalink
[BugFix] Wrong spec returned (#2604)
Browse files Browse the repository at this point in the history
  • Loading branch information
matteobettini authored Nov 24, 2024
1 parent 00d3199 commit a1e21f5
Showing 1 changed file with 4 additions and 3 deletions.
7 changes: 4 additions & 3 deletions torchrl/envs/common.py
Original file line number Diff line number Diff line change
Expand Up @@ -390,6 +390,7 @@ def __init__(
self.batch_size = torch.Size(batch_size)
self._run_type_checks = run_type_checks
self._allow_done_after_reset = allow_done_after_reset

@wraps(check_env_specs_func)
def check_env_specs(self, *args, **kwargs):
return check_env_specs_func(self, *args, **kwargs)
Expand Down Expand Up @@ -1550,7 +1551,7 @@ def action_spec_unbatched(self, spec: Composite):
@property
def full_observation_spec_unbatched(self) -> Composite:
"""Returns the observation spec of the env as if it had no batch dimensions."""
return self._make_single_env_spec(self.full_action_spec)
return self._make_single_env_spec(self.full_observation_spec)

@full_observation_spec_unbatched.setter
def full_observation_spec_unbatched(self, spec: Composite):
Expand All @@ -1570,7 +1571,7 @@ def observation_spec_unbatched(self, spec: Composite):
@property
def full_reward_spec_unbatched(self) -> Composite:
"""Returns the reward spec of the env as if it had no batch dimensions."""
return self._make_single_env_spec(self.full_action_spec)
return self._make_single_env_spec(self.full_reward_spec)

@full_reward_spec_unbatched.setter
def full_reward_spec_unbatched(self, spec: Composite):
Expand All @@ -1590,7 +1591,7 @@ def reward_spec_unbatched(self, spec: Composite):
@property
def full_done_spec_unbatched(self) -> Composite:
"""Returns the done spec of the env as if it had no batch dimensions."""
return self._make_single_env_spec(self.full_action_spec)
return self._make_single_env_spec(self.full_done_spec)

@full_done_spec_unbatched.setter
def full_done_spec_unbatched(self, spec: Composite):
Expand Down

0 comments on commit a1e21f5

Please sign in to comment.