tdmpc xarm example not working

[fredh-leap]

System Info

lerobot version=0.1.0
os=Ubuntu 20.04.6 LTS
python version=3.10.15
numpy version=1.26.4
torch version=2.4.1+cu121
lerobots configuration='dataset_repo_id': 'lerobot/xarm_lift_medium',
 'device': 'cuda',
 'env': {'action_dim': 4,
         'episode_length': 200,
         'fps': '${fps}',
         'gym': {'obs_type': 'pixels_agent_pos',
                 'render_mode': 'rgb_array',
                 'visualization_height': 384,
                 'visualization_width': 384},
         'image_size': 84,
         'name': 'xarm',
         'state_dim': 4,
         'task': 'XarmLift-v0'},
 'eval': {'batch_size': 1, 'n_episodes': 1, 'use_async_envs': False},
 'fps': 15,
 'policy': {'advantage_scaling': 3.0,
            'cem_iterations': 6,
            'consistency_coeff': 20.0,
            'discount': 0.9,
            'elite_weighting_temperature': 0.5,
            'expectile_weight': 0.9,
            'gaussian_mean_momentum': 0.1,
            'horizon': 5,
            'image_encoder_hidden_dim': 32,
            'input_normalization_modes': None,
            'input_shapes': {'observation.image': [3, 84, 84],
                             'observation.state': ['${env.state_dim}']},
            'latent_dim': 50,
            'max_random_shift_ratio': 0.0476,
            'max_std': 2.0,
            'min_std': 0.05,
            'mlp_dim': 512,
            'n_action_repeats': 2,
            'n_action_steps': 1,
            'n_elites': 50,
            'n_gaussian_samples': 512,
            'n_pi_samples': 51,
            'name': 'tdmpc',
            'output_normalization_modes': {'action': 'min_max'},
            'output_shapes': {'action': ['${env.action_dim}']},
            'pi_coeff': 0.5,
            'pretrained_model_path': None,
            'q_ensemble_size': 5,
            'reward_coeff': 0.5,
            'state_encoder_hidden_dim': 256,
            'target_model_momentum': 0.995,
            'temporal_decay_coeff': 0.5,
            'uncertainty_regularizer_coeff': 1.0,
            'use_mpc': True,
            'value_coeff': 0.1},
 'resume': False,
 'seed': 1,
 'training': {'batch_size': 256,
              'delta_timestamps': {'action': '[i / ${fps} for i in '
                                             'range(${policy.horizon})]',
                                   'next.reward': '[i / ${fps} for i in '
                                                  'range(${policy.horizon})]',
                                   'observation.image': '[i / ${fps} for i in '
                                                        'range(${policy.horizon} '
                                                        '+ 1)]',
                                   'observation.state': '[i / ${fps} for i in '
                                                        'range(${policy.horizon} '
                                                        '+ 1)]'},
              'do_online_rollout_async': False,
              'eval_freq': 5000,
              'grad_clip_norm': 10.0,
              'image_transforms': {'brightness': {'min_max': [0.8, 1.2],
                                                  'weight': 1},
                                   'contrast': {'min_max': [0.8, 1.2],
                                                'weight': 1},
                                   'enable': False,
                                   'hue': {'min_max': [-0.05, 0.05],
                                           'weight': 1},
                                   'max_num_transforms': 3,
                                   'random_order': False,
                                   'saturation': {'min_max': [0.5, 1.5],
                                                  'weight': 1},
                                   'sharpness': {'min_max': [0.8, 1.2],
                                                 'weight': 1}},
              'log_freq': 100,
              'lr': 0.0003,
              'num_workers': 4,
              'offline_steps': 50000,
              'online_buffer_capacity': 80000,
              'online_buffer_seed_size': 0,
              'online_env_seed': 10000,
              'online_rollout_batch_size': 1,
              'online_rollout_n_episodes': 1,
              'online_sampling_ratio': 0.5,
              'online_steps': 50000,
              'online_steps_between_rollouts': 50,
              'save_checkpoint': True,
              'save_freq': 10000},
 'use_amp': False,
 'video_backend': 'pyav',
 'wandb': {'disable_artifact': False,
           'enable': True,
           'notes': '',
           'project': 'lerobot'}}

Information

• One of the scripts in the examples/ folder of LeRobot
• My own task or dataset (give details below)

Reproduction

python lerobot/scripts/train.py dataset_repo_id=lerobot/xarm_lift_medium policy=tdmpc env=xarm device=cuda wandb.enable=true
INFO 2024-10-15 10:03:23 ts/train.py:244 {'dataset_repo_id': 'lerobot/xarm_lift_medium',
'device': 'cuda',
'env': {'action_dim': 4,
'episode_length': 200,
'fps': '${fps}',
'gym': {'obs_type': 'pixels_agent_pos',
'render_mode': 'rgb_array',
'visualization_height': 384,
'visualization_width': 384},
'image_size': 84,
'name': 'xarm',
'state_dim': 4,
'task': 'XarmLift-v0'},
'eval': {'batch_size': 1, 'n_episodes': 1, 'use_async_envs': False},
'fps': 15,
'policy': {'advantage_scaling': 3.0,
'cem_iterations': 6,
'consistency_coeff': 20.0,
'discount': 0.9,
'elite_weighting_temperature': 0.5,
'expectile_weight': 0.9,
'gaussian_mean_momentum': 0.1,
'horizon': 5,
'image_encoder_hidden_dim': 32,
'input_normalization_modes': None,
'input_shapes': {'observation.image': [3, 84, 84],
'observation.state': ['${env.state_dim}']},
'latent_dim': 50,
'max_random_shift_ratio': 0.0476,
'max_std': 2.0,
'min_std': 0.05,
'mlp_dim': 512,
'n_action_repeats': 2,
'n_action_steps': 1,
'n_elites': 50,
'n_gaussian_samples': 512,
'n_pi_samples': 51,
'name': 'tdmpc',
'output_normalization_modes': {'action': 'min_max'},
'output_shapes': {'action': ['${env.action_dim}']},
'pi_coeff': 0.5,
'pretrained_model_path': None,
'q_ensemble_size': 5,
'reward_coeff': 0.5,
'state_encoder_hidden_dim': 256,
'target_model_momentum': 0.995,
'temporal_decay_coeff': 0.5,
'uncertainty_regularizer_coeff': 1.0,
'use_mpc': True,
'value_coeff': 0.1},
'resume': False,
'seed': 1,
'training': {'batch_size': 256,
'delta_timestamps': {'action': '[i / ${fps} for i in '
'range(${policy.horizon})]',
'next.reward': '[i / ${fps} for i in '
'range(${policy.horizon})]',
'observation.image': '[i / ${fps} for i in '
'range(${policy.horizon} '
'+ 1)]',
'observation.state': '[i / ${fps} for i in '
'range(${policy.horizon} '
'+ 1)]'},
'do_online_rollout_async': False,
'eval_freq': 5000,
'grad_clip_norm': 10.0,
'image_transforms': {'brightness': {'min_max': [0.8, 1.2],
'weight': 1},
'contrast': {'min_max': [0.8, 1.2],
'weight': 1},
'enable': False,
'hue': {'min_max': [-0.05, 0.05],
'weight': 1},
'max_num_transforms': 3,
'random_order': False,
'saturation': {'min_max': [0.5, 1.5],
'weight': 1},
'sharpness': {'min_max': [0.8, 1.2],
'weight': 1}},
'log_freq': 100,
'lr': 0.0003,
'num_workers': 4,
'offline_steps': 50000,
'online_buffer_capacity': 80000,
'online_buffer_seed_size': 0,
'online_env_seed': 10000,
'online_rollout_batch_size': 1,
'online_rollout_n_episodes': 1,
'online_sampling_ratio': 0.5,
'online_steps': 50000,
'online_steps_between_rollouts': 50,
'save_checkpoint': True,
'save_freq': 10000},
'use_amp': False,
'video_backend': 'pyav',
'wandb': {'disable_artifact': False,
'enable': True,
'notes': '',
'project': 'lerobot'}}
Logs will be synced with wandb.
INFO 2024-10-15 10:03:25 n/logger.py:132 Track this run --> https://wandb.ai/frederick-j-hoffman-rare-event-ai/lerobot/runs/gevvjd8e
INFO 2024-10-15 10:03:25 ts/train.py:312 make_dataset
Fetching 806 files: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 806/806 [00:00<00:00, 9111.49it/s]
INFO 2024-10-15 10:03:28 ts/train.py:325 make_env
Error executing job with overrides: ['dataset_repo_id=lerobot/xarm_lift_medium', 'policy=tdmpc', 'env=xarm', 'device=cuda', 'wandb.enable=true']
Traceback (most recent call last):
File "/anaconda/envs/lerobot/lib/python3.10/site-packages/gymnasium/envs/registration.py", line 740, in make
env = env_creator(**env_spec_kwargs)
File "/anaconda/envs/lerobot/lib/python3.10/site-packages/gym_xarm/tasks/lift.py", line 16, in init
super().init("lift", **kwargs)
File "/anaconda/envs/lerobot/lib/python3.10/site-packages/gym_xarm/tasks/base.py", line 73, in init
self.observation_space = self._initialize_observation_space()
File "/anaconda/envs/lerobot/lib/python3.10/site-packages/gym_xarm/tasks/base.py", line 112, in _initialize_observation_space
obs = self.get_obs()
File "/anaconda/envs/lerobot/lib/python3.10/site-packages/gym_xarm/tasks/base.py", line 248, in get_obs
pixels = self._render(renderer=self.observation_renderer)
File "/anaconda/envs/lerobot/lib/python3.10/site-packages/gym_xarm/tasks/base.py", line 320, in _render
render = renderer.render(self.render_mode, camera_name="camera0")
TypeError: MujocoRenderer.render() got an unexpected keyword argument 'camera_name'

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
File "/mnt/batch/tasks/shared/LS_root/mounts/clusters/frednc41/code/Users/frederick.j.hoffman/lerobot/lerobot/scripts/train.py", line 652, in train_cli
train(
File "/mnt/batch/tasks/shared/LS_root/mounts/clusters/frednc41/code/Users/frederick.j.hoffman/lerobot/lerobot/scripts/train.py", line 326, in train
eval_env = make_env(cfg)
File "/mnt/batch/tasks/shared/LS_root/mounts/clusters/frednc41/code/Users/frederick.j.hoffman/lerobot/lerobot/common/envs/factory.py", line 51, in make_env
env = env_cls(
File "/anaconda/envs/lerobot/lib/python3.10/site-packages/gymnasium/vector/sync_vector_env.py", line 86, in init
self.envs = [env_fn() for env_fn in env_fns]
File "/anaconda/envs/lerobot/lib/python3.10/site-packages/gymnasium/vector/sync_vector_env.py", line 86, in
self.envs = [env_fn() for env_fn in env_fns]
File "/mnt/batch/tasks/shared/LS_root/mounts/clusters/frednc41/code/Users/frederick.j.hoffman/lerobot/lerobot/common/envs/factory.py", line 53, in
lambda: gym.make(gym_handle, disable_env_checker=True, **gym_kwgs)
File "/anaconda/envs/lerobot/lib/python3.10/site-packages/gymnasium/envs/registration.py", line 752, in make
raise type(e)(
TypeError: MujocoRenderer.render() got an unexpected keyword argument 'camera_name' was raised from the environment creator for gym_xarm/XarmLift-v0 with kwargs ({'obs_type': 'pixels_agent_pos', 'render_mode': 'rgb_array', 'visualization_width': 384, 'visualization_height': 384})

Set the environment variable HYDRA_FULL_ERROR=1 for a complete stack trace.

Expected behavior

I was hoping to train a simple tdmpc policy to get started