├── docs └── mujoco_sim.gif ├── conda-environment.yaml ├── LICENSE ├── src └── simulation │ ├── scene.xml │ ├── robot-02.xml │ ├── robot_lqr.py │ └── simulate_robot.py ├── README.md └── .gitignore /docs/mujoco_sim.gif: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lachlanhurst/balance-robot-mujoco-sim/HEAD/docs/mujoco_sim.gif -------------------------------------------------------------------------------- /conda-environment.yaml: -------------------------------------------------------------------------------- 1 | name: robot-mujoco-simulation 2 | channels: 3 | - conda-forge 4 | dependencies: 5 | - python==3.11 6 | - mujoco-python==3.2.0 7 | - pyside6==6.6.0 8 | - scipy==1.14.1 9 | - pip 10 | - pip: 11 | - control==0.10.1 12 | 13 | 14 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2024 Lachlan Hurst 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /src/simulation/scene.xml: -------------------------------------------------------------------------------- 1 | 2 | 3 | 36 | 37 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # Self balancing robot simulation in MuJoCo 2 | 3 | ## What / Why 4 | Simulation of a two wheeled self balancing robot, this time using [MuJoCo](https://mujoco.org/). The robot uses a Linear–quadratic regulator (LQR) controller to maintain balance. In this example the MuJoCo simulation view is embedded in a Qt application (PySide6) to provide UI widgets for controlling the robot. 5 | 6 | After implementing the same simulation using PyBullet (see [here](https://github.com/lachlanhurst/balance-robot-pybullet-sim)), I found that PyBullet support for [Stable-Baselines3](https://stable-baselines3.readthedocs.io/) / [Gymnasium](https://gymnasium.farama.org/index.html) was a little lacking. It made sense to switch to MuJoCo given how early on this project is. 7 | 8 | ![Self balancing robot simulation video](./docs/mujoco_sim.gif) 9 | 10 | 11 | ## Setting up the development environment 12 | 13 | Clone the repo 14 | 15 | git clone https://github.com/lachlanhurst/balance-robot-mujoco-sim.git 16 | cd balance-robot-mujoco-sim 17 | 18 | It's recommended to use the conda environment file included in this repository; these commands will create and activate a new `robot-mujoco-simulation` conda env. 19 | 20 | conda env create -f conda-environment.yaml 21 | conda activate robot-mujoco-simulation 22 | 23 | 24 | ## Calculating LQR parameters 25 | 26 | The LQR parameters are already [included in the source code](./src/simulation/robot_lqr.py). There is no need to re-calculate these parameters unless the physical characteristics of the robot change (as defined in [`robot-02.xml`](./src/simulation/robot-02.xml)). 27 | 28 | If you do need to re-calculate these parameters follow the process defined in the [PyBullet simulation repo](https://github.com/lachlanhurst/balance-robot-pybullet-sim). 29 | 30 | 31 | ## Running the simulation 32 | 33 | The following command will run the simulation. Note: the `robot-mujoco-simulation` must be activated 34 | 35 | python src/simulation/simulate_robot.py 36 | 37 | 38 | ## Acknowledgements 39 | 40 | [Jean Elsner](https://github.com/JeanElsner) put together this [great simple example](https://gist.github.com/JeanElsner/755d0feb49864ecadab4ef00fd49a22b) on how to embed a MuJoCo simulation within a PySide6 application. 41 | 42 | 43 | -------------------------------------------------------------------------------- /src/simulation/robot-02.xml: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | 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 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | # MuJoCo stuff 2 | MUJOCO_LOG.TXT 3 | 4 | 5 | 6 | ## Python stuff 7 | 8 | # Byte-compiled / optimized / DLL files 9 | __pycache__/ 10 | *.py[cod] 11 | *$py.class 12 | 13 | # C extensions 14 | *.so 15 | 16 | # Distribution / packaging 17 | .Python 18 | build/ 19 | develop-eggs/ 20 | dist/ 21 | downloads/ 22 | eggs/ 23 | .eggs/ 24 | lib/ 25 | lib64/ 26 | parts/ 27 | sdist/ 28 | var/ 29 | wheels/ 30 | share/python-wheels/ 31 | *.egg-info/ 32 | .installed.cfg 33 | *.egg 34 | MANIFEST 35 | 36 | # PyInstaller 37 | # Usually these files are written by a python script from a template 38 | # before PyInstaller builds the exe, so as to inject date/other infos into it. 39 | *.manifest 40 | *.spec 41 | 42 | # Installer logs 43 | pip-log.txt 44 | pip-delete-this-directory.txt 45 | 46 | # Unit test / coverage reports 47 | htmlcov/ 48 | .tox/ 49 | .nox/ 50 | .coverage 51 | .coverage.* 52 | .cache 53 | nosetests.xml 54 | coverage.xml 55 | *.cover 56 | *.py,cover 57 | .hypothesis/ 58 | .pytest_cache/ 59 | cover/ 60 | 61 | # Translations 62 | *.mo 63 | *.pot 64 | 65 | # Django stuff: 66 | *.log 67 | local_settings.py 68 | db.sqlite3 69 | db.sqlite3-journal 70 | 71 | # Flask stuff: 72 | instance/ 73 | .webassets-cache 74 | 75 | # Scrapy stuff: 76 | .scrapy 77 | 78 | # Sphinx documentation 79 | docs/_build/ 80 | 81 | # PyBuilder 82 | .pybuilder/ 83 | target/ 84 | 85 | # Jupyter Notebook 86 | .ipynb_checkpoints 87 | 88 | # IPython 89 | profile_default/ 90 | ipython_config.py 91 | 92 | # pyenv 93 | # For a library or package, you might want to ignore these files since the code is 94 | # intended to run in multiple environments; otherwise, check them in: 95 | # .python-version 96 | 97 | # pipenv 98 | # According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control. 99 | # However, in case of collaboration, if having platform-specific dependencies or dependencies 100 | # having no cross-platform support, pipenv may install dependencies that don't work, or not 101 | # install all needed dependencies. 102 | #Pipfile.lock 103 | 104 | # poetry 105 | # Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control. 106 | # This is especially recommended for binary packages to ensure reproducibility, and is more 107 | # commonly ignored for libraries. 108 | # https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control 109 | #poetry.lock 110 | 111 | # pdm 112 | # Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control. 113 | #pdm.lock 114 | # pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it 115 | # in version control. 116 | # https://pdm.fming.dev/latest/usage/project/#working-with-version-control 117 | .pdm.toml 118 | .pdm-python 119 | .pdm-build/ 120 | 121 | # PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm 122 | __pypackages__/ 123 | 124 | # Celery stuff 125 | celerybeat-schedule 126 | celerybeat.pid 127 | 128 | # SageMath parsed files 129 | *.sage.py 130 | 131 | # Environments 132 | .env 133 | .venv 134 | env/ 135 | venv/ 136 | ENV/ 137 | env.bak/ 138 | venv.bak/ 139 | 140 | # Spyder project settings 141 | .spyderproject 142 | .spyproject 143 | 144 | # Rope project settings 145 | .ropeproject 146 | 147 | # mkdocs documentation 148 | /site 149 | 150 | # mypy 151 | .mypy_cache/ 152 | .dmypy.json 153 | dmypy.json 154 | 155 | # Pyre type checker 156 | .pyre/ 157 | 158 | # pytype static type analyzer 159 | .pytype/ 160 | 161 | # Cython debug symbols 162 | cython_debug/ 163 | 164 | # PyCharm 165 | # JetBrains specific template is maintained in a separate JetBrains.gitignore that can 166 | # be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore 167 | # and can be added to the global gitignore or merged into this file. For a more nuclear 168 | # option (not recommended) you can uncomment the following to ignore the entire idea folder. 169 | #.idea/ 170 | -------------------------------------------------------------------------------- /src/simulation/robot_lqr.py: -------------------------------------------------------------------------------- 1 | import math 2 | 3 | import mujoco 4 | import numpy as np 5 | from scipy.spatial.transform import Rotation 6 | 7 | # obtained from running `calculate_lqr_gains.py` 8 | LQR_K = [-2.1402165848237837, -0.03501370844016172, 5.9748026764525894e-18, 2.236067977499789] 9 | WHEEL_RADIUS = 0.034 10 | MAX_MOTOR_VEL = 500.0 # rad/s 11 | 12 | 13 | def clamp(n, minn, maxn): 14 | return max(min(maxn, n), minn) 15 | 16 | 17 | """ 18 | Basic LQR controller implementation 19 | """ 20 | class RobotLqr: 21 | 22 | def __init__(self, model: mujoco.MjModel, data: mujoco.MjData) -> None: 23 | self.model = model 24 | self.data = data 25 | 26 | self.velocity_angular = 0.0 27 | self.velocity_linear_set_point = 0.0 28 | self.yaw = 0 29 | 30 | self.pitch_dot_filtered = 0.0 31 | self.velocity_angular_filtered = 0.0 32 | 33 | def set_velocity_linear_set_point(self, vel: float) -> None: 34 | """ 35 | Sets the target velocity of the robot 36 | """ 37 | self.velocity_linear_set_point = vel 38 | 39 | def set_yaw(self, yaw: float) -> None: 40 | """ 41 | The yaw value is added to one wheel, and subtracted from the other 42 | to produce a yaw motion (turn) 43 | """ 44 | self.yaw = yaw 45 | 46 | def get_pitch(self) -> float: 47 | quat = self.data.body("robot_body").xquat 48 | if quat[0] == 0: 49 | return 0 50 | 51 | rotation = Rotation.from_quat([quat[1], quat[2], quat[3], quat[0]]) # Quaternion order is [x, y, z, w] 52 | angles = rotation.as_euler('xyz', degrees=False) 53 | # print(angles) 54 | return angles[0] 55 | 56 | def get_pitch_dot(self) -> float: 57 | angular = self.data.joint('robot_body_joint').qvel[-3:] 58 | # print(angular) 59 | return angular[0] 60 | 61 | def get_wheel_velocity(self) -> float: 62 | vel_m_0 = self.data.joint('torso_l_wheel').qvel[0] 63 | vel_m_1 = self.data.joint('torso_r_wheel').qvel[0] 64 | 65 | # both wheels spin "forward", but one is spinning in a negative 66 | # direction as it's rotated 180deg from the other 67 | return (vel_m_0 * -1 + vel_m_1) / 2.0 68 | 69 | def calculate_lqr_velocity(self) -> float: 70 | pitch = -self.get_pitch() 71 | pitch_dot = self.get_pitch_dot() 72 | 73 | # apply a filter to pitch dot, and velocity 74 | # without these filters the controller seems to lack necessary dampening 75 | # would like to know why! 76 | self.pitch_dot_filtered = (self.pitch_dot_filtered * .975) + (pitch_dot * .025) 77 | self.velocity_angular_filtered = (self.velocity_angular_filtered * .975) + (self.get_wheel_velocity() * .025) 78 | 79 | velocity_linear_error = self.velocity_linear_set_point - self.velocity_angular_filtered * WHEEL_RADIUS 80 | 81 | lqr_v = LQR_K[0] * (0 - pitch) + LQR_K[1] * self.pitch_dot_filtered + LQR_K[2] * 0 + LQR_K[3] * velocity_linear_error 82 | return -lqr_v / WHEEL_RADIUS 83 | 84 | def update_motor_speed(self, ) -> None: 85 | vel = self.calculate_lqr_velocity() 86 | vel = clamp(vel, -MAX_MOTOR_VEL, MAX_MOTOR_VEL) 87 | 88 | self.data.actuator('motor_l_wheel').ctrl = [-vel + self.yaw] 89 | self.data.actuator('motor_r_wheel').ctrl = [vel + self.yaw] 90 | 91 | def reset(self): 92 | self.velocity_angular = 0.0 93 | self.velocity_linear_set_point = 0.0 94 | self.yaw = 0 95 | self.pitch_dot_filtered = 0.0 96 | self.velocity_angular_filtered = 0.0 97 | 98 | # face a random direction 99 | x_rot = (np.random.random() - 0.5) * 2 * math.pi 100 | # rotate and pitch slightly 101 | y_rot = (np.random.random() - 0.5) * 0.4 102 | z_rot = (np.random.random() - 0.5) * 0.4 103 | euler_angles = [x_rot, y_rot, z_rot] 104 | # Convert to quaternion 105 | rotation = Rotation.from_euler('xyz', euler_angles) 106 | self.data.qpos[3:7] = rotation.as_quat() 107 | 108 | self.data.actuator('motor_l_wheel').ctrl = [0] 109 | self.data.actuator('motor_r_wheel').ctrl = [0] 110 | -------------------------------------------------------------------------------- /src/simulation/simulate_robot.py: -------------------------------------------------------------------------------- 1 | from collections import deque 2 | import time 3 | import mujoco 4 | import numpy as np 5 | import pathlib 6 | from PySide6.QtWidgets import ( 7 | QApplication, QWidget, QMainWindow, QPushButton, QSizePolicy, 8 | QVBoxLayout, QGroupBox, QHBoxLayout, QSlider, QLabel 9 | ) 10 | from PySide6.QtCore import QTimer, Qt 11 | from PySide6.QtOpenGL import QOpenGLWindow 12 | from PySide6.QtWidgets import QApplication, QWidget 13 | from PySide6.QtCore import QTimer, Qt, Signal, Slot, QThread 14 | from PySide6.QtGui import ( 15 | QGuiApplication, QSurfaceFormat 16 | ) 17 | import time 18 | 19 | from robot_lqr import RobotLqr 20 | 21 | 22 | format = QSurfaceFormat() 23 | format.setDepthBufferSize(24) 24 | format.setStencilBufferSize(8) 25 | format.setSamples(4) 26 | format.setSwapInterval(1) 27 | format.setSwapBehavior(QSurfaceFormat.SwapBehavior.DoubleBuffer) 28 | format.setVersion(2,0) 29 | format.setRenderableType(QSurfaceFormat.RenderableType.OpenGL) 30 | format.setProfile(QSurfaceFormat.CompatibilityProfile) 31 | QSurfaceFormat.setDefaultFormat(format) 32 | 33 | 34 | class Viewport(QOpenGLWindow): 35 | 36 | updateRuntime = Signal(float) 37 | 38 | def __init__(self, model, data, cam, opt, scn) -> None: 39 | super().__init__() 40 | 41 | self.model = model 42 | self.data = data 43 | self.cam = cam 44 | self.opt = opt 45 | self.scn = scn 46 | 47 | self.width = 0 48 | self.height = 0 49 | self.scale = 1.0 50 | self.__last_pos = None 51 | 52 | self.runtime = deque(maxlen=1000) 53 | self.timer = QTimer() 54 | self.timer.setInterval(1/60*1000) 55 | self.timer.timeout.connect(self.update) 56 | self.timer.start() 57 | 58 | def mousePressEvent(self, event): 59 | self.__last_pos = event.position() 60 | 61 | def mouseMoveEvent(self, event): 62 | if event.buttons() & Qt.MouseButton.RightButton: 63 | action = mujoco.mjtMouse.mjMOUSE_MOVE_V 64 | elif event.buttons() & Qt.MouseButton.LeftButton: 65 | action = mujoco.mjtMouse.mjMOUSE_ROTATE_V 66 | elif event.buttons() & Qt.MouseButton.MiddleButton: 67 | action = mujoco.mjtMouse.mjMOUSE_ZOOM 68 | else: 69 | return 70 | pos = event.position() 71 | dx = pos.x() - self.__last_pos.x() 72 | dy = pos.y() - self.__last_pos.y() 73 | mujoco.mjv_moveCamera(self.model, action, dx / self.height, dy / self.height, self.scn, self.cam) 74 | self.__last_pos = pos 75 | 76 | def wheelEvent(self, event): 77 | mujoco.mjv_moveCamera(self.model, mujoco.mjtMouse.mjMOUSE_ZOOM, 0, -0.0005 * event.angleDelta().y(), self.scn, self.cam) 78 | 79 | def initializeGL(self): 80 | self.con = mujoco.MjrContext(self.model, mujoco.mjtFontScale.mjFONTSCALE_100) 81 | 82 | def resizeGL(self, w, h): 83 | self.width = w 84 | self.height = h 85 | 86 | def setScreenScale(self, scaleFactor: float) -> None: 87 | """ Sets a scale factor that is used to scale the OpenGL window to accommodate 88 | the high DPI scaling Qt does. 89 | """ 90 | self.scale = scaleFactor 91 | 92 | def paintGL(self) -> None: 93 | t = time.time() 94 | mujoco.mjv_updateScene(self.model, self.data, self.opt, None, self.cam, mujoco.mjtCatBit.mjCAT_ALL, self.scn) 95 | viewport = mujoco.MjrRect(0, 0, int(self.width * self.scale), int(self.height * self.scale)) 96 | mujoco.mjr_render(viewport, self.scn, self.con) 97 | 98 | self.runtime.append(time.time()-t) 99 | self.updateRuntime.emit(np.average(self.runtime)) 100 | 101 | 102 | class UpdateSimThread(QThread): 103 | 104 | def __init__(self, model: mujoco.MjModel, data: mujoco.MjData, parent=None) -> None: 105 | super().__init__(parent) 106 | self.model = model 107 | self.data = data 108 | self.running = True 109 | 110 | self.robot = RobotLqr(model, data) 111 | 112 | # robot control parameters 113 | self.speed = 0.0 114 | self.yaw = 0.0 115 | 116 | # reset the simulation timer 117 | self.reset() 118 | 119 | @property 120 | def real_time(self): 121 | return time.monotonic_ns() - self.real_time_start 122 | 123 | def run(self) -> None: 124 | while self.running: 125 | # don't step the simulation past real time 126 | # without this the sim usually finishes before it's 127 | # even visible 128 | if self.data.time < self.real_time / 1_000_000_000: 129 | # In the real robot we update the control loop at a 200hz, so do that 130 | # here too. It's the filters applied to pitch_dot and linear speed error 131 | # that are not time step independent 132 | if (time.monotonic_ns() - self.last_robot_update) / 1_000_000_000 >= (1/200): 133 | self.last_robot_update = time.monotonic_ns() 134 | # update robot with user inputs 135 | self.robot.set_velocity_linear_set_point(self.speed) 136 | self.robot.set_yaw(self.yaw) 137 | # update motor speed with LQR controller 138 | self.robot.update_motor_speed() 139 | 140 | # step the simulation 141 | mujoco.mj_step(self.model, self.data) 142 | else: 143 | time.sleep(0.00001) 144 | 145 | def stop(self): 146 | self.running = False 147 | self.wait() 148 | 149 | def reset(self): 150 | self.real_time_start = time.monotonic_ns() 151 | self.last_robot_update = time.monotonic_ns() 152 | self.robot.reset() 153 | 154 | def set_speed(self, speed: float) -> None: 155 | self.speed = speed 156 | 157 | def set_yaw(self, yaw: float) -> None: 158 | self.yaw = yaw 159 | 160 | 161 | class Window(QMainWindow): 162 | 163 | def __init__(self) -> None: 164 | super().__init__() 165 | 166 | self.model = mujoco.MjModel.from_xml_path(str(pathlib.Path(__file__).parent.joinpath('scene.xml'))) 167 | self.data = mujoco.MjData(self.model) 168 | self.cam = self.create_free_camera() 169 | self.opt = mujoco.MjvOption() 170 | self.scn = mujoco.MjvScene(self.model, maxgeom=10000) 171 | self.scn.flags[mujoco.mjtRndFlag.mjRND_SHADOW] = True 172 | self.scn.flags[mujoco.mjtRndFlag.mjRND_REFLECTION] = True 173 | self.viewport = Viewport(self.model, self.data, self.cam, self.opt, self.scn) 174 | self.viewport.setScreenScale(QGuiApplication.instance().primaryScreen().devicePixelRatio()) 175 | self.viewport.updateRuntime.connect(self.show_runtime) 176 | 177 | layout = QVBoxLayout() 178 | layout_top = QHBoxLayout() 179 | layout_top.setSpacing(8) 180 | reset_button = QPushButton("Reset") 181 | reset_button.setMinimumWidth(90) 182 | reset_button.setSizePolicy(QSizePolicy.Fixed, QSizePolicy.Expanding) 183 | reset_button.clicked.connect(self.reset_simulation) 184 | layout_top.addWidget(reset_button) 185 | layout_robot_controls = QVBoxLayout() 186 | layout_robot_controls.setContentsMargins(0,0,0,0) 187 | layout_robot_controls.addWidget(self.create_top()) 188 | layout_top.addLayout(layout_robot_controls) 189 | layout_top.setContentsMargins(8,0,8,0) 190 | layout.addLayout(layout_top) 191 | layout.addWidget(QWidget.createWindowContainer(self.viewport)) 192 | layout.setContentsMargins(0,4,0,0) 193 | layout.setStretch(1,1) 194 | w = QWidget() 195 | w.setLayout(layout) 196 | self.setCentralWidget(w) 197 | self.resize(800, 600) 198 | 199 | self.th = UpdateSimThread(self.model, self.data, self) 200 | self.th.start() 201 | 202 | @Slot(float) 203 | def show_runtime(self, fps: float): 204 | self.statusBar().showMessage( 205 | f"Average runtime: {fps:.0e}s\t" 206 | f"Simulation time: {self.data.time:.0f}s" 207 | ) 208 | 209 | def create_top(self): 210 | layout = QVBoxLayout() 211 | # layout.setContentsMargins(0,0,0,0) 212 | label_width = 60 213 | 214 | speed_layout = QHBoxLayout() 215 | self.speed_slider = QSlider(Qt.Horizontal) 216 | # QSliders only support ints, so scale the values we want by 1000 217 | # and then remove this scale factor in the valueChanged handler 218 | self.speed_slider.setMinimum(-4 * 1000) 219 | self.speed_slider.setMaximum(4 * 1000) 220 | self.speed_slider.setValue(0) 221 | self.speed_slider.valueChanged.connect(self._speed_changed) 222 | speed_label = QLabel("Speed") 223 | speed_label.setFixedWidth(label_width) 224 | speed_layout.addWidget(speed_label) 225 | speed_layout.addWidget(self.speed_slider) 226 | 227 | yaw_layout = QHBoxLayout() 228 | self.yaw_slider = QSlider(Qt.Horizontal) 229 | self.yaw_slider.setMinimum(-10 * 1000) 230 | self.yaw_slider.setMaximum(10 * 1000) 231 | self.yaw_slider.setValue(0) 232 | self.yaw_slider.valueChanged.connect(self._yaw_changed) 233 | yaw_label = QLabel("Yaw") 234 | yaw_label.setFixedWidth(label_width) 235 | yaw_layout.addWidget(yaw_label) 236 | yaw_layout.addWidget(self.yaw_slider) 237 | 238 | layout.addLayout(speed_layout) 239 | layout.addLayout(yaw_layout) 240 | 241 | w = QGroupBox("Robot Control") 242 | w.setLayout(layout) 243 | return w 244 | 245 | def _speed_changed(self, value: int) -> None: 246 | speed = value / 1000 247 | self.th.set_speed(speed) 248 | 249 | def _yaw_changed(self, value: int) -> None: 250 | yaw = value / 1000 251 | self.th.set_yaw(yaw) 252 | 253 | def create_free_camera(self): 254 | cam = mujoco.MjvCamera() 255 | cam.type = mujoco.mjtCamera.mjCAMERA_FREE 256 | cam.fixedcamid = -1 257 | cam.lookat = np.array([ 0.0 , 0.0 , 0.0 ]) 258 | cam.distance = self.model.stat.extent * 2 259 | cam.elevation = -25 260 | cam.azimuth = 45 261 | return cam 262 | 263 | def reset_simulation(self): 264 | self.speed_slider.setValue(0) 265 | self.yaw_slider.setValue(0) 266 | # Reset state and time. 267 | mujoco.mj_resetData(self.model, self.data) 268 | mujoco.mj_forward(self.model, self.data) 269 | self.th.reset() 270 | 271 | 272 | if __name__ == "__main__": 273 | app = QApplication() 274 | w = Window() 275 | w.show() 276 | app.exec() 277 | w.th.stop() 278 | --------------------------------------------------------------------------------