|
22 | 22 | R = TypeVar("R") |
23 | 23 |
|
24 | 24 |
|
| 25 | +class Control(StrEnum): |
| 26 | + """Control type of the simulated onboard controller.""" |
| 27 | + |
| 28 | + state = "state" |
| 29 | + """State control takes [x, y, z, vx, vy, vz, ax, ay, az, qx, qy, qz, qw, wx, wy, wz]. |
| 30 | +
|
| 31 | + Note: |
| 32 | + Recommended frequency is >=20 Hz. |
| 33 | +
|
| 34 | + Warning: |
| 35 | + Only the yaw of the attitude quaternion is used, as in the firmware. The so_rpy family |
| 36 | + ignores the body rate setpoint. |
| 37 | + """ |
| 38 | + attitude = "attitude" |
| 39 | + """Attitude control takes [roll, pitch, yaw, collective thrust]. |
| 40 | +
|
| 41 | + Note: |
| 42 | + Recommended frequency is >=100 Hz. |
| 43 | + """ |
| 44 | + body_rate = "body_rate" |
| 45 | + """Body rate control takes [wx, wy, wz, collective thrust]. |
| 46 | +
|
| 47 | + Note: |
| 48 | + Recommended frequency is >=200 Hz. |
| 49 | + """ |
| 50 | + force_torque = "force_torque" |
| 51 | + """Force and torque control takes [fc, tx, ty, tz]. |
| 52 | +
|
| 53 | + Note: |
| 54 | + Recommended frequency is >=500 Hz. |
| 55 | + """ |
| 56 | + rotor_vel = "rotor_vel" |
| 57 | + """Rotor velocity control takes [w1, w2, w3, w4] in RPMs. |
| 58 | +
|
| 59 | + Note: |
| 60 | + Recommended frequency is >=500 Hz. |
| 61 | + """ |
| 62 | + default = attitude |
| 63 | + |
| 64 | + |
25 | 65 | def parametrize( |
26 | 66 | fn: Callable[P, R], drone: str, xp: ModuleType | None = None, device: str | None = None |
27 | 67 | ) -> Callable[P, R]: |
@@ -50,77 +90,57 @@ def parametrize( |
50 | 90 | Returns: |
51 | 91 | The parametrized controller function with all keyword argument only parameters filled in. |
52 | 92 | """ |
53 | | - return _parametrize(fn, drone, load_params, xp=xp, device=device) |
| 93 | + return _parametrize(fn, drone, load_fn_params, xp=xp, device=device) |
54 | 94 |
|
55 | 95 |
|
56 | 96 | def load_params( |
57 | | - fn: Callable, drone: str, xp: ModuleType | None = None, device: str | None = None |
58 | | -) -> dict[str, Array]: |
59 | | - """Load the parameters a specific controller function accepts. |
| 97 | + controller: str, drone: str, xp: ModuleType | None = None, device: str | None = None |
| 98 | +) -> dict[str, dict[str, Array]]: |
| 99 | + """Load all parameters of a drone for a controller. |
60 | 100 |
|
61 | | - Merges the ``"core"`` section with the function's ``[drone.<fn_name>]`` section (function values |
62 | | - take precedence), then keeps only the parameters in ``fn``'s signature. |
| 101 | + Returns the drone's table of ``crazyflow/control/<controller>/params.toml``: the ``core`` |
| 102 | + section shared by all functions of the controller, and one section per controller function. |
63 | 103 |
|
64 | 104 | Args: |
65 | | - fn: The controller function for which to load parameters. |
| 105 | + controller: Name of the controller package, e.g. ``"mellinger"``. |
66 | 106 | drone: Name of the drone configuration, e.g. ``"cf2x_L250"``. |
67 | 107 | xp: The array API module to use. If not provided, numpy is used. |
68 | 108 | device: The device to use. If None, the device is inferred from the xp module. |
69 | 109 |
|
70 | | - Returns: |
71 | | - A flat dict mapping parameter names to arrays in the requested array namespace. |
| 110 | + Raises: |
| 111 | + KeyError: If ``controller`` has no ``params.toml`` or ``drone`` has no section in it. |
72 | 112 | """ |
73 | | - assert isinstance(fn, Callable), f"Expected a function, got {type(fn)}" |
74 | | - controller = fn.__module__.split(".")[-2] |
75 | 113 | params_path = Path(__file__).parent / f"{controller}/params.toml" |
76 | 114 | if not params_path.exists(): |
77 | | - raise KeyError(f"`{controller}` not found. Available controllers: {tuple(Control)}") |
| 115 | + raise KeyError(f"Controller `{controller}` not found") |
78 | 116 | with open(params_path, "rb") as f: |
79 | 117 | params = tomllib.load(f) |
80 | 118 | if drone not in params: |
81 | 119 | raise KeyError(f"Drone `{drone}` not found in {controller}/params.toml") |
82 | | - merged = params[drone].get("core", {}) | params[drone].get(fn.__name__, {}) |
83 | | - return to_xp(filter_to_signature(merged, fn), xp=xp, device=device) |
84 | | - |
85 | | - |
86 | | -class Control(StrEnum): |
87 | | - """Control type of the simulated onboard controller.""" |
88 | | - |
89 | | - state = "state" |
90 | | - """State control takes [x, y, z, vx, vy, vz, ax, ay, az, qx, qy, qz, qw, wx, wy, wz]. |
91 | | -
|
92 | | - Note: |
93 | | - Recommended frequency is >=20 Hz. |
| 120 | + return to_xp(params[drone], xp=xp, device=device) |
94 | 121 |
|
95 | | - Warning: |
96 | | - Only the yaw of the attitude quaternion is used, as in the firmware. The so_rpy family |
97 | | - ignores the body rate setpoint. |
98 | | - """ |
99 | | - attitude = "attitude" |
100 | | - """Attitude control takes [roll, pitch, yaw, collective thrust]. |
101 | 122 |
|
102 | | - Note: |
103 | | - Recommended frequency is >=100 Hz. |
104 | | - """ |
105 | | - body_rate = "body_rate" |
106 | | - """Body rate control takes [wx, wy, wz, collective thrust]. |
| 123 | +def load_fn_params( |
| 124 | + fn: Callable, drone: str, xp: ModuleType | None = None, device: str | None = None |
| 125 | +) -> dict[str, Array]: |
| 126 | + """Load the parameters a controller function accepts. |
107 | 127 |
|
108 | | - Note: |
109 | | - Recommended frequency is >=200 Hz. |
110 | | - """ |
111 | | - force_torque = "force_torque" |
112 | | - """Force and torque control takes [fc, tx, ty, tz]. |
| 128 | + The controller is derived from the function's package. Merges the ``core`` section with the |
| 129 | + function's section (function values take precedence), then keeps only the parameters in |
| 130 | + ``fn``'s signature. |
113 | 131 |
|
114 | | - Note: |
115 | | - Recommended frequency is >=500 Hz. |
116 | | - """ |
117 | | - rotor_vel = "rotor_vel" |
118 | | - """Rotor velocity control takes [w1, w2, w3, w4] in RPMs. |
| 132 | + Args: |
| 133 | + fn: The controller function for which to load parameters. |
| 134 | + drone: Name of the drone configuration, e.g. ``"cf2x_L250"``. |
| 135 | + xp: The array API module to use. If not provided, numpy is used. |
| 136 | + device: The device to use. If None, the device is inferred from the xp module. |
119 | 137 |
|
120 | | - Note: |
121 | | - Recommended frequency is >=500 Hz. |
| 138 | + Returns: |
| 139 | + A flat dict mapping parameter names to arrays in the requested array namespace. |
122 | 140 | """ |
123 | | - default = attitude |
| 141 | + assert callable(fn), f"Expected a function, got {type(fn)}" |
| 142 | + params = load_params(fn.__module__.split(".")[-2], drone, xp=xp, device=device) |
| 143 | + return filter_to_signature(params.get("core", {}) | params.get(fn.__name__, {}), fn) |
124 | 144 |
|
125 | 145 |
|
126 | 146 | @jax.jit |
|
0 commit comments