Add todo for train

This commit is contained in:
Simon Alibert 2024-05-05 15:04:23 +02:00
parent 17778fa3cb
commit 41acc73b4a
2 changed files with 3 additions and 2 deletions

View File

@ -152,8 +152,9 @@ See `python lerobot/scripts/eval.py --help` for more instructions.
Check out [example 3](./examples/3_train_policy.py) to see how you can start training a model on a dataset, which will be automatically downloaded if needed. Check out [example 3](./examples/3_train_policy.py) to see how you can start training a model on a dataset, which will be automatically downloaded if needed.
In general, you can use our training script to easily train any policy in any environment: In general, you can use our training script to easily train any policy on its environment:
```bash ```bash
# TODO(aliberts): not working
python lerobot/scripts/train.py \ python lerobot/scripts/train.py \
env=aloha \ env=aloha \
task=sim_insertion \ task=sim_insertion \

View File

@ -18,7 +18,7 @@ from lerobot.common.utils.utils import init_hydra_config
output_directory = Path("outputs/train/example_pusht_diffusion") output_directory = Path("outputs/train/example_pusht_diffusion")
os.makedirs(output_directory, exist_ok=True) os.makedirs(output_directory, exist_ok=True)
# Number of offline training steps (we'll only do offline training for this example. # Number of offline training steps (we'll only do offline training for this example.)
# Adjust as you prefer. 5000 steps are needed to get something worth evaluating. # Adjust as you prefer. 5000 steps are needed to get something worth evaluating.
training_steps = 5000 training_steps = 5000
device = torch.device("cuda") device = torch.device("cuda")