Skip to content

[BugFix] Parametric collectors #1303

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.

Sign up for GitHub

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

Merged
merged 14 commits into from
Jun 29, 2023
Merged

[BugFix] Parametric collectors #1303

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
14 commits
Select commit Hold shift + click to select a range
e5f5d89
parametric collectors
matteobettini Jun 21, 2023
afc18b5
revert
matteobettini Jun 21, 2023
f6e4c3d
revert
matteobettini Jun 21, 2023
8c27cff
fix
matteobettini Jun 21, 2023
d4c95a0
added test
matteobettini Jun 22, 2023
8d5fc9f
new tests
matteobettini Jun 22, 2023
4b9fdb4
lint
matteobettini Jun 22, 2023
8ed9d69
empty
vmoens Jun 22, 2023
4958d8d
fix and simplify test
matteobettini Jun 23, 2023
05cf021
slimmer tests
matteobettini Jun 28, 2023
2f721c5
Merge branch 'main' into parametric_collectors
matteobettini Jun 28, 2023
7c39c89
empty
matteobettini Jun 28, 2023
290d8a5
empty
vmoens Jun 22, 2023
d834e33
Merge remote-tracking branch 'origin/main' into parametric_collectors
vmoens Jun 28, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
169 changes: 146 additions & 23 deletions test/mocking_classes.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,13 +7,16 @@
import torch
import torch.nn as nn
from tensordict.tensordict import TensorDict, TensorDictBase
from tensordict.utils import NestedKey

from torchrl.data.tensor_specs import (
BinaryDiscreteTensorSpec,
BoundedTensorSpec,
CompositeSpec,
DiscreteTensorSpec,
MultiOneHotDiscreteTensorSpec,
OneHotDiscreteTensorSpec,
TensorSpec,
UnboundedContinuousTensorSpec,
)
from torchrl.envs.common import EnvBase
Expand Down Expand Up @@ -941,6 +944,15 @@ def forward(self, observation, action):
return self.linear(torch.cat([observation, action], dim=-1))


class CountingEnvCountPolicy:
def __init__(self, action_spec: TensorSpec, action_key: NestedKey = "action"):
self.action_spec = action_spec
self.action_key = action_key

def __call__(self, td: TensorDictBase) -> TensorDictBase:
return td.set(self.action_key, self.action_spec.zero() + 1)


class CountingEnv(EnvBase):
"""An env that is done after a given number of steps.

Expand Down Expand Up @@ -1011,7 +1023,7 @@ def _step(
self,
tensordict: TensorDictBase,
) -> TensorDictBase:
action = tensordict.get("action")
action = tensordict.get(self.action_key)
self.count += action.to(torch.int).to(self.device)
tensordict = TensorDict(
source={
Expand All @@ -1025,38 +1037,149 @@ def _step(
return tensordict.select().set("next", tensordict)


class NestedRewardEnv(CountingEnv):
class NestedCountingEnv(CountingEnv):
# an env with nested reward and done states
def __init__(self, max_steps: int = 5, start_val: int = 0, **kwargs):
def __init__(
self,
max_steps: int = 5,
start_val: int = 0,
nest_obs_action: bool = True,
nest_done: bool = True,
nest_reward: bool = True,
nested_dim: int = 3,
**kwargs,
):
super().__init__(max_steps=max_steps, start_val=start_val, **kwargs)
self.observation_spec = CompositeSpec(
{("data", "states"): self.observation_spec["observation"].clone()},
shape=self.batch_size,
)
self.reward_spec = CompositeSpec(
{("data", "reward"): self.reward_spec.clone()}, shape=self.batch_size
)
self.done_spec = CompositeSpec(
{("data", "done"): self.done_spec.clone()}, shape=self.batch_size
)

self.nested_dim = nested_dim

self.nested_obs_action = nest_obs_action
self.nested_done = nest_done
self.nested_reward = nest_reward

if self.nested_obs_action:
self.observation_spec = CompositeSpec(
{
"data": CompositeSpec(
{
"states": self.observation_spec["observation"]
.unsqueeze(-1)
.expand(*self.batch_size, self.nested_dim, 1)
},
shape=(
*self.batch_size,
self.nested_dim,
),
)
},
shape=self.batch_size,
)
self.action_spec = CompositeSpec(
{
"data": CompositeSpec(
{
"action": self.action_spec.unsqueeze(-1).expand(
*self.batch_size, self.nested_dim, 1
)
},
shape=(
*self.batch_size,
self.nested_dim,
),
)
},
shape=self.batch_size,
)

if self.nested_reward:
self.reward_spec = CompositeSpec(
{
"data": CompositeSpec(
{
"reward": self.reward_spec.unsqueeze(-1).expand(
*self.batch_size, self.nested_dim, 1
)
},
shape=(
*self.batch_size,
self.nested_dim,
),
)
},
shape=self.batch_size,
)

if self.nested_done:
self.done_spec = CompositeSpec(
{
"data": CompositeSpec(
{
"done": self.done_spec.unsqueeze(-1).expand(
*self.batch_size, self.nested_dim, 1
)
},
shape=(
*self.batch_size,
self.nested_dim,
),
)
},
shape=self.batch_size,
)

def _reset(self, td):
if self.nested_done and td is not None and "_reset" in td.keys():
td["_reset"] = td["_reset"].sum(-2, dtype=torch.bool)
td = super()._reset(td)
td[self.done_key] = td["done"]
del td["done"]
td["data", "states"] = td["observation"]
del td["observation"]
if self.nested_done:
td[self.done_key] = (
td["done"].unsqueeze(-1).expand(*self.batch_size, self.nested_dim, 1)
)
del td["done"]
if self.nested_obs_action:
td["data", "states"] = (
td["observation"]
.unsqueeze(-1)
.expand(*self.batch_size, self.nested_dim, 1)
)
del td["observation"]
if "data" in td.keys():
td["data"].batch_size = (*self.batch_size, self.nested_dim)
return td

def _step(self, td):
if self.nested_obs_action:
td["data"].batch_size = self.batch_size
td[self.action_key] = td[self.action_key].max(-2)[0]
td_root = super()._step(td)
if self.nested_obs_action:
td[self.action_key] = (
td[self.action_key]
.unsqueeze(-1)
.expand(*self.batch_size, self.nested_dim, 1)
)
if "data" in td.keys():
td["data"].batch_size = (*self.batch_size, self.nested_dim)
td = td_root["next"]
td[self.reward_key] = td["reward"]
del td["reward"]
td[self.done_key] = td["done"]
del td["done"]
td["data", "states"] = td["observation"]
del td["observation"]
if self.nested_done:
td[self.done_key] = (
td["done"].unsqueeze(-1).expand(*self.batch_size, self.nested_dim, 1)
)
del td["done"]
if self.nested_obs_action:
td["data", "states"] = (
td["observation"]
.unsqueeze(-1)
.expand(*self.batch_size, self.nested_dim, 1)
)
del td["observation"]
if self.nested_reward:
td[self.reward_key] = (
td["reward"].unsqueeze(-1).expand(*self.batch_size, self.nested_dim, 1)
)
del td["reward"]
if "data" in td.keys():
td["data"].batch_size = (*self.batch_size, self.nested_dim)
return td_root


Expand Down
126 changes: 124 additions & 2 deletions test/test_collector.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,16 +14,18 @@
ContinuousActionVecMockEnv,
CountingBatchedEnv,
CountingEnv,
CountingEnvCountPolicy,
DiscreteActionConvMockEnv,
DiscreteActionConvPolicy,
DiscreteActionVecMockEnv,
DiscreteActionVecPolicy,
MockSerialEnv,
NestedCountingEnv,
)
from tensordict.nn import TensorDictModule
from tensordict.tensordict import assert_allclose_td, TensorDict
from torch import nn
from torchrl._utils import seed_generator
from torchrl._utils import prod, seed_generator
from torchrl.collectors import aSyncDataCollector, SyncDataCollector
from torchrl.collectors.collectors import (
_Interruptor,
Expand All @@ -33,7 +35,14 @@
)
from torchrl.collectors.utils import split_trajectories
from torchrl.data import CompositeSpec, UnboundedContinuousTensorSpec
from torchrl.envs import EnvBase, EnvCreator, ParallelEnv, SerialEnv, StepCounter
from torchrl.envs import (
EnvBase,
EnvCreator,
InitTracker,
ParallelEnv,
SerialEnv,
StepCounter,
)
from torchrl.envs.libs.gym import _has_gym, GymEnv
from torchrl.envs.transforms import TransformedEnv, VecNorm
from torchrl.modules import Actor, LSTMNet, OrnsteinUhlenbeckProcessWrapper, SafeModule
Expand Down Expand Up @@ -1346,6 +1355,119 @@ def test_reset_heterogeneous_envs():
).all()


class TestNestedEnvsCollector:
def test_multi_collector_nested_env_consistency(self, seed=1):
env = NestedCountingEnv()
torch.manual_seed(seed)
env_fn = lambda: TransformedEnv(env, InitTracker())
policy = CountingEnvCountPolicy(env.action_spec, env.action_key)

ccollector = MultiaSyncDataCollector(
create_env_fn=[env_fn],
policy=policy,
frames_per_batch=20,
total_frames=100,
device="cpu",
)
for i, d in enumerate(ccollector):
if i == 0:
c1 = d
elif i == 1:
c2 = d
else:
break
assert d.names[-1] == "time"
with pytest.raises(AssertionError):
assert_allclose_td(c1, c2)
ccollector.shutdown()

ccollector = MultiSyncDataCollector(
create_env_fn=[env_fn],
policy=policy,
frames_per_batch=20,
total_frames=100,
device="cpu",
)
for i, d in enumerate(ccollector):
if i == 0:
d1 = d
elif i == 1:
d2 = d
else:
break
assert d.names[-1] == "time"
with pytest.raises(AssertionError):
assert_allclose_td(d1, d2)
ccollector.shutdown()

assert_allclose_td(c1, d1)
assert_allclose_td(c2, d2)

@pytest.mark.parametrize("nested_obs_action", [True, False])
@pytest.mark.parametrize("nested_done", [True, False])
@pytest.mark.parametrize("nested_reward", [True, False])
def test_collector_nested_env_combinations(
self,
nested_obs_action,
nested_done,
nested_reward,
seed=1,
frames_per_batch=20,
):
env = NestedCountingEnv(
nest_reward=nested_reward,
nest_done=nested_done,
nest_obs_action=nested_obs_action,
)
torch.manual_seed(seed)
policy = CountingEnvCountPolicy(env.action_spec, env.action_key)
ccollector = SyncDataCollector(
create_env_fn=env,
policy=policy,
frames_per_batch=frames_per_batch,
total_frames=100,
device="cpu",
)

for _td in ccollector:
break
ccollector.shutdown()

@pytest.mark.parametrize("batch_size", [(), (5,), (5, 2)])
def test_nested_env_dims(self, batch_size, nested_dim=5, frames_per_batch=20):
from mocking_classes import CountingEnvCountPolicy, NestedCountingEnv

env = NestedCountingEnv(batch_size=batch_size, nested_dim=nested_dim)
env_fn = lambda: NestedCountingEnv(batch_size=batch_size, nested_dim=nested_dim)
torch.manual_seed(0)
policy = CountingEnvCountPolicy(env.action_spec, env.action_key)
policy(env.reset())
ccollector = SyncDataCollector(
create_env_fn=env_fn,
policy=policy,
frames_per_batch=frames_per_batch,
total_frames=100,
device="cpu",
)

for _td in ccollector:
break
ccollector.shutdown()

# assert ("data","reward") not in td.keys(True) # this can be activates once step_mdp is fixed for nested keys
Copy link
Contributor Author

@matteobettini matteobettini Jun 22, 2023 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

this needs #1268 fixed to be uncommented

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Got it, afaic we can merge it and uncomment later

assert _td.batch_size == (*batch_size, frames_per_batch // prod(batch_size))
assert _td["data"].batch_size == (
*batch_size,
frames_per_batch // prod(batch_size),
nested_dim,
)
assert _td["next", "data"].batch_size == (
*batch_size,
frames_per_batch // prod(batch_size),
nested_dim,
)


@pytest.mark.skipif(not torch.cuda.device_count(), reason="No casting if no cuda")
class TestUpdateParams:
class DummyEnv(EnvBase):
Expand Down
Loading