judo/judo/controller/overrides.py at 1ec93f9e17808629ea702fe1b837a3ae6985de65 · bhung-bdai/judo · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
# Copyright (c) 2025 Robotics and AI Institute LLC. All rights reserved.

from judo.config import set_config_overrides
from judo.controller.controller import ControllerConfig


def set_default_cylinder_push_overrides() -> None:
    """Sets the default task-specific controller config overrides for the cylinder push task."""
    set_config_overrides(
        "cylinder_push",
        ControllerConfig,
        {
            "horizon": 1.0,
            "spline_order": "zero",
        },
    )


def set_default_cartpole_overrides() -> None:
    """Sets the default task-specific controller config overrides for the cartpole task."""
    set_config_overrides(
        "cartpole",
        ControllerConfig,
        {
            "horizon": 1.0,
            "spline_order": "zero",
        },
    )


def set_default_leap_cube_overrides() -> None:
    """Sets the default task-specific controller config overrides for the leap cube task."""
    set_config_overrides(
        "leap_cube",
        ControllerConfig,
        {
            "horizon": 1.0,
            "spline_order": "cubic",
            "max_num_traces": 1,
        },
    )


def set_default_leap_cube_down_overrides() -> None:
    """Sets the default task-specific controller config overrides for the leap cube down task."""
    set_config_overrides(
        "leap_cube_down",
        ControllerConfig,
        {
            "horizon": 1.0,
            "spline_order": "cubic",
            "max_num_traces": 1,
        },
    )


def set_default_caltech_leap_cube_overrides() -> None:
    """Sets the default task-specific controller config overrides for the caltech leap cube task."""
    set_config_overrides(
        "caltech_leap_cube",
        ControllerConfig,
        {
            "horizon": 1.0,
            "spline_order": "cubic",
            "max_num_traces": 1,
        },
    )


_SPOT_TASK_NAMES = [
    "spot_base",
    "spot_box_push",
    "spot_navigate",
    "spot_tire_roll",
    "spot_tire_upright",
]


def set_default_spot_overrides() -> None:
    """Sets the default task-specific controller config overrides for all Spot tasks."""
    for task_name in _SPOT_TASK_NAMES:
        set_config_overrides(
            task_name,
            ControllerConfig,
            {
                "horizon": 2.0,
            },
        )


def set_default_fr3_pick_overrides() -> None:
    """Sets the default task-specific controller config overrides for the fr3 pick task."""
    set_config_overrides(
        "fr3_pick",
        ControllerConfig,
        {
            "horizon": 1.0,
            "spline_order": "linear",
            "max_num_traces": 3,
            "control_freq": 20.0,
        },
    )