How to train the original Robomimic task square?

[dzy666fly]

Thank u so much for sharing this good work!

I have run the code successfully with the command in the README.md. Well, I am also interested in the Robomimic task such as Square, Can, so I want to test EquiDiff method on these tasks. But if I simply add the file robomimic_abs.yaml which mimics the mimicgen_abs.yaml in the code, it doesen't work and shows the error as following. So how can I do that?

(equidiff) zhenyuan:~/Code/Source/equidiff$ CUDA_VISIBLE_DEVICES=1 MUJOCO_GL=osmesa PYOPENGL_PLATTFORM=osmesa python train.py --config-name=robomimic_train_act_abs task_name=square n_demo=180 training.debug=True
Error executing job with overrides: ['task_name=square', 'n_demo=180', 'training.debug=True']
Traceback (most recent call last):
  File "/home/zhenyuan/Code/Source/equidiff/train.py", line 67, in <module>
    main()
  File "/home/zhenyuan/Softwares/miniforge3/envs/equidiff/lib/python3.9/site-packages/hydra/main.py", line 90, in decorated_main
    _run_hydra(
  File "/home/zhenyuan/Softwares/miniforge3/envs/equidiff/lib/python3.9/site-packages/hydra/_internal/utils.py", line 389, in _run_hydra
    _run_app(
  File "/home/zhenyuan/Softwares/miniforge3/envs/equidiff/lib/python3.9/site-packages/hydra/_internal/utils.py", line 452, in _run_app
    run_and_report(
  File "/home/zhenyuan/Softwares/miniforge3/envs/equidiff/lib/python3.9/site-packages/hydra/_internal/utils.py", line 216, in run_and_report
    raise ex
  File "/home/zhenyuan/Softwares/miniforge3/envs/equidiff/lib/python3.9/site-packages/hydra/_internal/utils.py", line 213, in run_and_report
    return func()
  File "/home/zhenyuan/Softwares/miniforge3/envs/equidiff/lib/python3.9/site-packages/hydra/_internal/utils.py", line 453, in <lambda>
    lambda: hydra.run(
  File "/home/zhenyuan/Softwares/miniforge3/envs/equidiff/lib/python3.9/site-packages/hydra/_internal/hydra.py", line 132, in run
    _ = ret.return_value
  File "/home/zhenyuan/Softwares/miniforge3/envs/equidiff/lib/python3.9/site-packages/hydra/core/utils.py", line 260, in return_value
    raise self._return_value
  File "/home/zhenyuan/Softwares/miniforge3/envs/equidiff/lib/python3.9/site-packages/hydra/core/utils.py", line 186, in run_job
    ret.return_value = task_function(task_cfg)
  File "/home/zhenyuan/Code/Source/equidiff/train.py", line 60, in main
    OmegaConf.resolve(cfg)
  File "/home/zhenyuan/Softwares/miniforge3/envs/equidiff/lib/python3.9/site-packages/omegaconf/omegaconf.py", line 771, in resolve
    omegaconf._impl._resolve(cfg)
  File "/home/zhenyuan/Softwares/miniforge3/envs/equidiff/lib/python3.9/site-packages/omegaconf/_impl.py", line 40, in _resolve
    _resolve_container_value(cfg, k)
  File "/home/zhenyuan/Softwares/miniforge3/envs/equidiff/lib/python3.9/site-packages/omegaconf/_impl.py", line 14, in _resolve_container_value
    resolved = node._dereference_node()
  File "/home/zhenyuan/Softwares/miniforge3/envs/equidiff/lib/python3.9/site-packages/omegaconf/base.py", line 246, in _dereference_node
    node = self._dereference_node_impl(throw_on_resolution_failure=True)
  File "/home/zhenyuan/Softwares/miniforge3/envs/equidiff/lib/python3.9/site-packages/omegaconf/base.py", line 277, in _dereference_node_impl
    return parent._resolve_interpolation_from_parse_tree(
  File "/home/zhenyuan/Softwares/miniforge3/envs/equidiff/lib/python3.9/site-packages/omegaconf/base.py", line 584, in _resolve_interpolation_from_parse_tree
    resolved = self.resolve_parse_tree(
  File "/home/zhenyuan/Softwares/miniforge3/envs/equidiff/lib/python3.9/site-packages/omegaconf/base.py", line 764, in resolve_parse_tree
    return visitor.visit(parse_tree)
  File "/home/zhenyuan/Softwares/miniforge3/envs/equidiff/lib/python3.9/site-packages/antlr4/tree/Tree.py", line 34, in visit
    return tree.accept(self)
  File "/home/zhenyuan/Softwares/miniforge3/envs/equidiff/lib/python3.9/site-packages/omegaconf/grammar/gen/OmegaConfGrammarParser.py", line 206, in accept
    return visitor.visitConfigValue(self)
  File "/home/zhenyuan/Softwares/miniforge3/envs/equidiff/lib/python3.9/site-packages/omegaconf/grammar_visitor.py", line 101, in visitConfigValue
    return self.visit(ctx.getChild(0))
  File "/home/zhenyuan/Softwares/miniforge3/envs/equidiff/lib/python3.9/site-packages/antlr4/tree/Tree.py", line 34, in visit
    return tree.accept(self)
  File "/home/zhenyuan/Softwares/miniforge3/envs/equidiff/lib/python3.9/site-packages/omegaconf/grammar/gen/OmegaConfGrammarParser.py", line 342, in accept
    return visitor.visitText(self)
  File "/home/zhenyuan/Softwares/miniforge3/envs/equidiff/lib/python3.9/site-packages/omegaconf/grammar_visitor.py", line 301, in visitText
    return self._unescape(list(ctx.getChildren()))
  File "/home/zhenyuan/Softwares/miniforge3/envs/equidiff/lib/python3.9/site-packages/omegaconf/grammar_visitor.py", line 389, in _unescape
    text = str(self.visitInterpolation(node))
  File "/home/zhenyuan/Softwares/miniforge3/envs/equidiff/lib/python3.9/site-packages/omegaconf/grammar_visitor.py", line 125, in visitInterpolation
    return self.visit(ctx.getChild(0))
  File "/home/zhenyuan/Softwares/miniforge3/envs/equidiff/lib/python3.9/site-packages/antlr4/tree/Tree.py", line 34, in visit
    return tree.accept(self)
  File "/home/zhenyuan/Softwares/miniforge3/envs/equidiff/lib/python3.9/site-packages/omegaconf/grammar/gen/OmegaConfGrammarParser.py", line 921, in accept
    return visitor.visitInterpolationNode(self)
  File "/home/zhenyuan/Softwares/miniforge3/envs/equidiff/lib/python3.9/site-packages/omegaconf/grammar_visitor.py", line 158, in visitInterpolationNode
    return self.node_interpolation_callback(inter_key, self.memo)
  File "/home/zhenyuan/Softwares/miniforge3/envs/equidiff/lib/python3.9/site-packages/omegaconf/base.py", line 745, in node_interpolation_callback
    return self._resolve_node_interpolation(inter_key=inter_key, memo=memo)
  File "/home/zhenyuan/Softwares/miniforge3/envs/equidiff/lib/python3.9/site-packages/omegaconf/base.py", line 676, in _resolve_node_interpolation
    raise InterpolationKeyError(f"Interpolation key '{inter_key}' not found")
omegaconf.errors.InterpolationKeyError: Interpolation key 'task.task_name' not found

Looking forward to ur reply!

[dzy666fly]

Well, after add the robomimic dataset, it seems to work. But I am not sure the robomimic environment of EquiDiff matches the original Robomimic tasks.

[pointW]

Hi @dzy666fly, I would recommend re-generating the visual observation in the robomimic dataset. The original robomimic used an older version of mujoco, so the visual observation would be slightly different from the test environment if you use the robomimic dataset directly.