mirror of
https://github.com/huggingface/lerobot.git
synced 2026-05-16 09:09:48 +00:00
Add memory optimization option to ReplayBuffer
- Introduce `optimize_memory` parameter to reduce memory usage in replay buffer - Implement simplified memory optimization by not storing duplicate next_states - Update learner server and buffer initialization to use memory optimization by default
This commit is contained in:
@@ -154,6 +154,7 @@ def initialize_replay_buffer(
|
||||
device=device,
|
||||
state_keys=cfg.policy.input_shapes.keys(),
|
||||
storage_device=device,
|
||||
optimize_memory=True,
|
||||
)
|
||||
|
||||
dataset = LeRobotDataset(
|
||||
@@ -166,6 +167,7 @@ def initialize_replay_buffer(
|
||||
capacity=cfg.training.online_buffer_capacity,
|
||||
device=device,
|
||||
state_keys=cfg.policy.input_shapes.keys(),
|
||||
optimize_memory=True,
|
||||
)
|
||||
|
||||
|
||||
@@ -648,6 +650,7 @@ def train(cfg: DictConfig, out_dir: str | None = None, job_name: str | None = No
|
||||
action_mask=active_action_dims,
|
||||
action_delta=cfg.env.wrapper.delta_action,
|
||||
storage_device=device,
|
||||
optimize_memory=True,
|
||||
)
|
||||
batch_size: int = batch_size // 2 # We will sample from both replay buffer
|
||||
|
||||
|
||||
Reference in New Issue
Block a user