hangX/lerobot
1
0
Fork 0
forked from tangger/lerobot
Code Pull Requests Activity

Merge remote-tracking branch 'refs/remotes/origin/2024_05_30_add_data_augmentation' into 2024_05_30_add_data_augmentation

Browse Source
This commit is contained in:
Marina Barannikov
2024-06-04 12:00:46 +00:00
parent 5eea2542d9 7be2c35c0a
commit 31e3c82386
4 changed files with 7 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
lerobot/common/policies/diffusion/modeling_diffusion.py
View File

@@ -239,10 +239,8 @@ class DiffusionModel(nn.Module):
global_cond = torch.cat([batch["observation.state"], img_features], dim=-1).flatten(start_dim=1)
# run sampling
sample = self.conditional_sample(batch_size, global_cond=global_cond)
actions = self.conditional_sample(batch_size, global_cond=global_cond)
# `horizon` steps worth of actions (from the first observation).
actions = sample[..., : self.config.output_shapes["action"][0]]
# Extract `n_action_steps` steps worth of actions (from the current observation).
start = n_obs_steps - 1
end = start + self.config.n_action_steps

2
lerobot/common/policies/normalize.py
View File

@@ -147,7 +147,7 @@ class Normalize(nn.Module):
assert not torch.isinf(min).any(), _no_stats_error_str("min")
assert not torch.isinf(max).any(), _no_stats_error_str("max")
# normalize to [0,1]
batch[key] = (batch[key] - min) / (max - min)
batch[key] = (batch[key] - min) / (max - min + 1e-8)
# normalize to [-1, 1]
batch[key] = batch[key] * 2 - 1
else: