# mjlab_gs

Canonical standalone simulator workspace for the UR5/RG2 bottle scene, with
MJLab tasks and calibrated Gaussian wrist RGB. BC, RL, testing, recording and
evaluation all construct `ManagerBasedRlEnv` from the same task registration.

```bash
cd ~/mjlab_gs
source env.sh
mjlab-gs list
mjlab-gs smoke --num-envs 4 --steps 20
mjlab-gs view --num-envs 4
```

Open http://127.0.0.1:8080. The viewer starts paused. Its **Play, Step, Reset,
environment selector, camera feeds and debugging panels** are MJLab's Viser
viewer, built on mjviser. The interactive 3D view shows the meshes; the wrist RGB
feed shows the actual hybrid Gaussian image supplied to the policy. Browser
streaming copies selected images/state to the CPU; headless training does not
use the viewer.

For a browser on your laptop, forward the GPU machine's port:

```bash
ssh -L 8080:127.0.0.1:8080 <your-GPU-host>
# Then open http://127.0.0.1:8080 on the laptop.
```

`--host` and `--port` are configurable. `--running` starts playback immediately.
`--checkpoint` works with canonical BC and PPO checkpoints.

## Workspace

- `mjlab/`: independent upstream clone, pinned at
  `8ee51fbcf806a7419189f706d9e394cbeb7790fa` (MJLab 1.6.0), with local patches.
- `src/mjlab_gs/tasks/`: scene composition, GPU actions, reset distribution,
  observations, curriculum, domain randomization, rewards and task registration.
- `src/mjlab_gs/rendering/`: calibrated gsplat sensor and corrected Warp renderer.
- `assets/`, `assets_lod/`, `scene/`: copied checkpoints, CAD, meshes, textures,
  calibration and scene XML. No symlinks to another project.
- `assets-manifest.json`: hashes for every copied asset and calibration file.
- `.venv/`, `toolchains/`, `torch_extensions/`: local Python dependencies, copied
  CUDA compilation tools, and locally rebuilt gsplat extension.
- `scripts/`, `tests/`, `reports/`: verification and measurements. The copied
  CPU reference under `tests/reference/` is only a validation aid.

There is no runtime dependency on GSim, rl_training, policy_training or real2sim
folders. System NVIDIA drivers and a Python interpreter are still required.
After moving the workspace or rebuilding the environment, run
`scripts/setup.sh` and source `env.sh`. First execution can compile kernels for
several minutes. The large checkpoint/asset files are already included.

## Shared interface

Task: **`Mjlab-GS-Bottle-UR5`**. Default 64 environments, 320×200 RGB, 20 Hz
control, 50 physics substeps at 1 ms. Camera width/height are configurable;
320×200 preserves the original benchmark resolution. Packed SH3 gsplat rendering
uses camera chunks of 128, shadowed Warp foregrounds, and validated visual LOD.
Collision meshes and physical parameters remain those of the source scene.

```python
from mjlab_gs import make_env
import torch

env = make_env(num_envs=64, play=True, auto_reset=False)
obs, info = env.reset()
action = torch.zeros((64, 7), device=env.device)
action[:, 6] = 1  # open gripper
obs, reward, terminated, truncated, info = env.step(action)
rgb_uint8 = env.scene['wrist'].data.rgb  # CUDA [N,H,W,3]
env.close()
```

The `mjlab-gs-cartesian-v3` contract action has 7 normalized values:
`[dx,dy,dz,rx,ry,rz,gripper]`. Translation and axis-angle rotation are expressed
in the robot base frame, with vector norms limited to 5 mm and 1 degree per
control step. MJLab's GPU differential IK runs each physics substep, with bounded-actuator
bias compensation to hold the arm against gravity. Gripper
`-1` means closed and `+1` means open; a zero gripper action requests half-open.
The utility hold policy explicitly requests open.

Observations: `actor` = 16-value proprioception (six joint angles, normalized
pad separation, base-frame TCP position and two rotation columns); `camera` =
CUDA float32 `[N,3,H,W]` in [0,1]; `goal` adds 3 robot-base target coordinates; `critic` has 62 privileged state values.
Raw RGB is uint8 for recording. Checkpoints contain a simulator contract and
reject incompatible configurations on load. Original CPU-IK BC checkpoints
are **not behaviorally interchangeable**: converting units alone does not
reproduce that controller's iterative IK, interpolation or gravity compensation.

Training uses staged domain randomization for object/table/robot poses, initial
joints, camera calibration, physics and appearance. Gaussian pedestal motion and
native rendering share per-world transforms. Play/record/evaluation are nominal
by default; use `--randomize`, `--dr-scale 1` or `--task-config configs/curriculum.json`.
See [the complete task specification](task.md) for ranges, schedule,
observation ordering, rewards, success criteria, metrics and limitations.

## Workflows

Normal MJLab entry points work after `source env.sh`:

```bash
list-envs
play Mjlab-GS-Bottle-UR5 --agent zero --viewer viser
train Mjlab-GS-Bottle-UR5 --env.scene.num-envs 64
```

Convenience entry points use the same registered task:

```bash
mjlab-gs train --num-envs 64 --iterations 3000 --output logs/rl/bottle
mjlab-gs evaluate --num-envs 8 --steps 300 --checkpoint logs/rl/bottle/final.pt
mjlab-gs record --num-envs 4 --steps 100 --output logs/data/rollout.npz
mjlab-gs bc-train logs/data/rollout.npz --updates 100 --output logs/bc/model.pt
mjlab-gs evaluate --checkpoint logs/bc/model.pt --num-envs 4
mjlab-gs benchmark --num-envs 64 --steps 100 --output reports/benchmark64.json
```

The task directly reuses MJLab lift-cube position shaping, action smoothness,
joint-limit and joint-speed penalties, including the native penalty curriculum.
A reset-only goal adapter samples reachable points around the nominal TCP,
independently of bottle placement. Success requires a valid bilateral
10-cm lift held for 2 seconds. Evaluation reports fixed-quota success/lift metrics
and tables by DR scale and bottle location. Episodes continue after success to
reward sustained holding. The visual actor receives 16 proprioceptive values + 3 goal coordinates + RGB; the critic is 62 values
including bottle-to-goal displacement. See [goal and teacher update](teacher-report.md).
No trained grasp policy is included.
`record` without a checkpoint produces hold-open plumbing transitions, **not
expert demonstrations**. The simple BC trainer is for small datasets that fit
GPU memory. Recordings include terminal observations before manual resets;
disk recording intentionally transfers to the CPU. Production datasets need a
streaming loader.

PPO uses the MJLab visual actor and a privileged state critic, TensorBoard logs,
normalization, checkpointing and resume-compatible environment counters. Its
rollout buffer holds float32 camera observations: memory grows with environment
count × rollout length × image resolution. Simulator-only capacity is not PPO
training capacity.

## GPU execution and measurements

Physics, Cartesian IK, camera transforms, foreground rendering, Gaussian
rasterization, compositing, rewards and observations operate on CUDA. Initialization
loads/compiles assets on the CPU. Packed gsplat still retrieves small allocation
counts from the GPU, and MJLab may synchronize episode/reset bookkeeping. This
is not a completely host-free pipeline. The IK solve avoids a CPU status check
on each substep using `solve_ex(check_errors=False)` on the damped system.

The benchmark command measures the **whole MJLab control step**, including IK,
50 physics steps, reward/reset/observation managers and RGB conversion. It excludes
policy inference/optimization and viewer/disk I/O. The old GSim held-control
physics benchmark is a lower-cost workload and its FPS is not a promise for this
task. Torch memory counters exclude Warp allocations; consult NVML/process memory
alongside them for total GPU usage.

Validation commands:

```bash
python -m pytest tests/test_contract.py -q
python scripts/assets_manifest.py --verify
python scripts/check_model.py
CUDA_VISIBLE_DEVICES=0 python scripts/validate.py
```

## Local MJLab changes

Review with `git -C mjlab diff`:

1. `SimulationCfg.render_fn`: optional renderer hook, default preserves upstream.
2. `SensorContext.recreate`: rebind sensors after rebuilding render context.
3. `EntityCfg.mocap_fixed_base`: optional permanently welded fixed entities;
   default preserves upstream mocap wrapping.
4. `DifferentialIKActionCfg.check_solve_errors`: optional asynchronous solve;
   default preserves upstream error checking.
5. Viser camera elevation: convert MuJoCo's negative-above convention to
   mjviser's positive-above convention for the initial browser view.

The robot task, Gaussian composition and workflow code remain outside the clone.
See `THIRD_PARTY.md` for source and license provenance.

The canonical IK substep is captured as a CUDA graph to reduce Python/operator
launch overhead; `cfg.actions['arm'].capture_ik=False` provides the uncaptured
reference path. Model field expansion triggers recapture. `env.sh` enables
`PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True`, which prevents changing gsplat
visibility counts from accumulating large cached allocation blocks as the camera
moves. Pre-reset observation shape probes return a zero image until the first
reset initializes the camera pose.


## Rendering-free state teacher

`Mjlab-GS-Bottle-UR5-State` is the low-dimensional teacher variant: 37 state
observations, actor/critic MLPs 256–128–64, no camera initialization, and its own
physical DR defaults in `configs/teacher.json`. The visual task now includes an
explicit 3D goal input; both use reachable randomized nominal-TCP goals and the
same control/reward/metric implementation. Observation contract v3 rejects old
non-goal-conditioned actors. See [teacher/student workflow](teacher.md).

```bash
mjlab-gs view --task Mjlab-GS-Bottle-UR5-State --num-envs 4
mjlab-gs benchmark --task Mjlab-GS-Bottle-UR5-State --num-envs 256 --task-config configs/teacher.json
```


## Current visual default

Wrist images default to **320×200**, including the task configuration, CLI and
matrix evaluator. The visual PPO actor uses ResNet18 + an ACT-style 256-wide,
8-head, 3-encoder/2-decoder transformer, then a 256-unit action head. Visual BC
shares its feature architecture. The teacher remains a 256–128–64 state MLP.
The visual actor is goal-conditioned and produces one action per 20-Hz tick;
it does not yet recreate ACT's 16-frame history or chunked action execution.
See [current visual checks](visual-report.md). Prior benchmark
numbers are tied to the image dimensions recorded in each run.