tangger/lerobot
1
0
Fork 1
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

Add storage device parameter to replay buffer initialization

Browse Source 
- Specify storage device for replay buffer to optimize memory management
This commit is contained in:
AdilZouitine
2025-02-25 15:30:39 +00:00
committed by Michel Aractingi
parent 1d4ec50a58
commit 9ea79f8a76
1 changed files with 1 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
lerobot/scripts/server/learner_server.py
View File

@@ -647,6 +647,7 @@ def train(cfg: DictConfig, out_dir: str | None = None, job_name: str | None = No
state_keys=cfg.policy.input_shapes.keys(), state_keys=cfg.policy.input_shapes.keys(),
action_mask=active_action_dims, action_mask=active_action_dims,
action_delta=cfg.env.wrapper.delta_action, action_delta=cfg.env.wrapper.delta_action,
storage_device=device,
) )
batch_size: int = batch_size // 2 # We will sample from both replay buffer batch_size: int = batch_size // 2 # We will sample from both replay buffer