VLA-MuJoco-Learning

Imported from xkang-zhao/VLA-MuJoco-Learning (AGENTS.md). Install upstream with npx skills add xkang-zhao/VLA-MuJoco-Learning. Copyright stays with the author.

AGENTS.md

This file provides guidance to Codex (Codex.ai/code) when working with code in this repository.

Project Overview

MuJoCo simulation of a UR10e robot arm with 2F85 gripper for reinforcement learning research. The robot is floor-mounted (ground-fixed base), controlled via keyboard teleoperation through inverse kinematics (IK).

Dev Environment & Commands

所有操作均在 conda 虚拟环境 myrobot 中进行。

conda activate myrobot

# 安装依赖（仅初次）
pip install -r requirements.txt
pip install -e .

Run the simulation with keyboard control (Mac):

python scripts/keyboard_control.py

Run multi-task keyboard teleoperation:

python scripts/keyboard_control.py --task cube

Run unit tests:

python test/test_controller.py
python test/test_kinematics.py

Testing Rule

每次新增功能、代码或修复 bug 时，必须在 test/ 目录下添加对应的测试代码，并运行测试确保通过。 测试应覆盖核心逻辑，不要求覆盖需要 GUI/硬件（如 MuJoCo viewer）的集成场景。

运行单个测试：

python test/<test_file>.py

Reference Repository

my_simulation — 本地路径：./reference_repos/my_simulation/

该仓库用作代码参考（如 IK 求解、控制器逻辑、MJCF 模型结构等）。仅作参考，绝对不能修改或提交到该仓库。

Architecture

Package Layers (bottom-up)

1. src/mujoco_robot/ — Core simulation library (pip install -e . installs this as package mujoco_robot)

   • Kinematics (robot_kinematics.py): Mink-based FK/IK solver (MuJoCo 原生微分 IK). 使用 QP + Levenberg-Marquardt 阻尼求解，支持关节限位和 warm start。Pose 表示 [x, y, z, roll, pitch, yaw] (ZYX Euler).
   • RobotController (robot_controller.py): Maintains target pose state, applies tool-frame delta updates via 4×4 transform right-multiplication, manages data.ctrl and data.qpos.
   • RobotSensor (robot_sensor.py): MuJoCo renderer wrapper for RGB/depth camera rendering.
   • RobotViewer (robot_viewer.py): Passive MuJoCo viewer wrapper.

2. src/env/gym_env.py — UR10eEnv(gym.Env): Gymnasium environment. Wires together Kinematics, Controller, Sensor, Viewer. Action space is 7D normalized [-1,1]: [dx, dy, dz, droll, dpitch, dyaw, gripper]. Observation includes joint positions (6), EE pose (7), target cube pose (7), and camera images (640×480 RGB + optional depth). Reward is negative distance to target cube.

3. src/env/stack_env.py — UR10eStackEnv: Red-blue cube stacking task. Extends UR10eEnv with success detection, randomized spawns, and dense/sparse rewards.

4. src/env/insert_env.py — UR10eInsertEnv: Free-peg vertical insertion task. Configures contact properties, peg damping, and success criteria.

MJCF Model Files

• mjcf/ur10e.xml — UR10e arm definition (includes mesh assets from mjcf/assets/)
• mjcf/ur10e_2f85.xml — Combines UR10e arm + 2F85 gripper
• mjcf/scene.xml — Top-level scene: includes the robot, plus floor, cube (free body with red box geom), world-frame visualization, and top_camera. Timestep: 0.0005s, implicitfast integrator.

Key Conventions

• Pose representation throughout: [x, y, z, roll, pitch, yaw] (ZYX Euler angles). Not quaternions (except in Gym obs where EE/cube pose includes xyzw quat).
• Tool-frame incremental control: delta actions are right-multiplied as 4×4 transforms onto the current target pose (T_new = T_current @ T_delta).
• MuJoCo joint ordering: qpos[0:6] = arm joints, ctrl[0:6] = arm position commands, ctrl[6] = gripper command (0–255 range).