# Linear Quadratic Regulator LQR (Linear Quadratic Regulator) is a classic continuous control and stabilization task. This repository currently provides two variants: - `dm-lqr-2-1`: two masses connected by a rope, with only the last mass actuated - `dm-lqr-6-2`: six masses connected as a chain, with only the last two masses actuated The goal is to drive the whole system back to the center and keep it near equilibrium with minimal control effort. ```{video} /_static/videos/dm_lqr_2_1.mp4 :poster: _static/images/poster/dm_lqr_2_1.jpg :nocontrols: :autoplay: :playsinline: :muted: :loop: :width: 100% ``` ```{video} /_static/videos/dm_lqr_6_2.mp4 :poster: _static/images/poster/dm_lqr_6_2.jpg :nocontrols: :autoplay: :playsinline: :muted: :loop: :width: 100% ``` --- ## Task Description Both tasks can be viewed as one-dimensional spring-damper chain stabilization problems. Each mass has a single translational degree of freedom along the x-axis. Neighboring masses are coupled by rope-like spring forces, and the system is affected by: - body damping on each mass - spring forces and relative damping between neighboring masses - a center-restoring force pulling the system toward the origin - control inputs applied only to the actuated terminal degrees of freedom In practice: - `dm-lqr-2-1` is the simpler version and is useful for verifying whether the policy can learn a stable equilibrium - `dm-lqr-6-2` is more difficult because the controller must propagate its effect through a longer chain --- ## Action Space ### dm-lqr-2-1 | Item | Details | | ------------- | ------------------------------- | | **Type** | `Box(-1.0, 1.0, (1,), float32)` | | **Dimension** | 1 | | Index | Action Description | Min | Max | XML Joint | | ----- | -------------------------------------- | ---- | --- | --------- | | 0 | Control input applied to the last mass | -1.0 | 1.0 | `q1` | ### dm-lqr-6-2 | Item | Details | | ------------- | ------------------------------- | | **Type** | `Box(-1.0, 1.0, (2,), float32)` | | **Dimension** | 2 | | Index | Action Description | Min | Max | XML Joint | | ----- | --------------------------------------------- | ---- | --- | --------- | | 0 | Control input applied to the second-last mass | -1.0 | 1.0 | `q4` | | 1 | Control input applied to the last mass | -1.0 | 1.0 | `q5` | --- ## Observation Space The observation is formed by concatenating all positions `qpos` and velocities `qvel`. ### dm-lqr-2-1 | Item | Details | | ------------- | ------------------------------- | | **Type** | `Box(-inf, inf, (4,), float32)` | | **Dimension** | 4 | | Index | Observation | Meaning | | ----- | ----------- | --------------------------- | | 0 | `q0` | Position of the first mass | | 1 | `q1` | Position of the second mass | | 2 | `dq0` | Velocity of the first mass | | 3 | `dq1` | Velocity of the second mass | ### dm-lqr-6-2 | Item | Details | | ------------- | -------------------------------- | | **Type** | `Box(-inf, inf, (12,), float32)` | | **Dimension** | 12 | The first 6 dimensions are `q0 ~ q5`, and the last 6 dimensions are `dq0 ~ dq5`. --- ## Reward Function Design The current reward is composed of state cost, velocity cost, control cost, success bonus, and out-of-bounds penalty: ```python state_cost = 0.5 * sum(qpos ** 2) velocity_cost = 0.5 * velocity_cost_coef * sum(qvel ** 2) control_cost = 0.5 * control_cost_coef * sum(action ** 2) reward = 1.0 - (state_cost + velocity_cost + control_cost) reward += success_bonus reward -= out_of_bounds_penalty ``` Intuitively: - the farther the system is from the origin, the lower the reward - larger velocities reduce the reward - aggressive control inputs reduce the reward - entering a small stable region around the origin yields a success bonus - leaving the valid state boundary triggers an additional penalty --- ## Initial State At reset: - the position vector is sampled in a random direction and normalized to a fixed norm - all initial velocities are set to zero With the current configuration: - `dm-lqr-2-1` starts with position norm around `0.8` - `dm-lqr-6-2` starts with position norm around `1.0` --- ## Episode Termination Conditions An episode terminates and resets when any of the following conditions is met: - success condition is reached: the position norm is below the success distance threshold and the velocity norm is below the success velocity threshold - out-of-bounds condition is reached: any position exceeds the position boundary or any velocity exceeds the velocity boundary - the full state is sufficiently close to zero - `NaN` appears in the observation or action --- ## Usage Guide ### 1. Environment Preview ```bash uv run scripts/view.py --env dm-lqr-2-1 uv run scripts/view.py --env dm-lqr-6-2 ``` ### 2. Start Training ```bash uv run scripts/train.py --env dm-lqr-2-1 uv run scripts/train.py --env dm-lqr-6-2 ``` ### 3. View Training Progress ```bash uv run tensorboard --logdir runs/dm-lqr-2-1 uv run tensorboard --logdir runs/dm-lqr-6-2 ``` ### 4. Test Training Results ```bash uv run scripts/play.py --env dm-lqr-2-1 uv run scripts/play.py --env dm-lqr-6-2 ``` --- ## Expected Training Results ### dm-lqr-2-1 1. The actuated mass pulls the unactuated mass back toward the center. 2. Both positions and velocities converge to a small neighborhood of zero. 3. The learned policy does not settle at a biased off-center equilibrium. ### dm-lqr-6-2 1. The last two actuated masses gradually pull the entire chain back toward the center. 2. The chain remains stable without obvious divergence or persistent oscillation. 3. Success rate increases during training while the out-of-bounds rate decreases.