Support for Gym-like environments

[yardenas]

First I would like to say thanks for this great work!

I am currently working on running mjswan within my slidev presentation. It works very nicely. It's pretty cool to demo policies trained on ogbench directly in my slides.

However, while developing this feature, I had to modify mjswan (specifically some parts in runtime.ts etc.).

It seems like currently mjswan is geared towards mjlab-style "mdp" environments, which is cool, however somewhat limited.

It seems very simple these days to just ask codex/claude code to scan a python RL environment and translate it to ts (my codex agent zero-shotted cube and puzzle tasks from ogbench).

I was wondering if the authors would be interested in a contribution that allows users to plug-in gym-like environments? My understanding is that the core changes are in runtime.ts?

[ttktjmt]

Thank you for the positive feedback and the suggestion!

I'd like to keep the current compatibility with mjlab, but I'm also interested in adding an adapter for gym-like environments if feasible.

Could you share a bit more detail about what kinds of changes would be needed to introduce this feature, and what role the adapter would play in making it work? It would help me get a clearer picture of the scope and design before moving forward.

[yardenas]

I envision something that roughly follows the design suggested below.

Environment adapter

• Essentially follow a gym-like abstraction:

reset(options?: { seed?: number }): void;
observe(): Float32Array;
step(action: Float32Array): Promise<EnvStepInfo> | EnvStepInfo;
info(): EnvStepInfo;
visualUpdates?(): EnvVisualUpdate[];  // Environment specific visual changes to the environment, e.g. "button changed color from blue to red"

• Abstraction by mjswan that allows registering new RL environments with their assets. Specifically, in loadEnvironment in runtime.ts, implement await this.loadEnvAdapter(scenePath), which will get all the assets and the TS implementation of the RL environment.

Policy

Simpler setup than what exists now:

• Instead of specifying how to read/transform observations from MuJoCo, feed directly the observation to the onnx policy. Currently this is done by this.policyRunner.collectObservationsByKey), which I believe we won't need for this path.
• This means it's the user's responsibility to implement their environment observation in a way that matches the original python environment.

Simulation loop

Add a path in mainLoop that:

0. calls reset()
1. calls visualUpdates
2. applies external user forces
3. runs onnx policy (in the future we could abstract this to any policy, like an oracle policy)
4. calls step(action), emits info for more user info rendering (like visualizing the ongoing reward)

Open questions

• Since we'll step the environment from runtime.ts, this implementation will require users to reimplement their RL (usually python) environments in TS. I believe with coding agents this is not very hard: codex could zero-shot tasks from ogbench quite directly. Still, if we could just load an existing python implementation (via pyodide?) that would significantly increase usability of this method.

I'm happy to share my code for this demo if that further helps

Demo

[ttktjmt]

Yes, I want to take a look at your codebase

[yardenas]

I added you to the repo!

[yardenas]

Repo is here: https://github.com/yardenas/mjswan-gym

[yardenas]

Feel free to reach out at yardas@ethz.ch if something is unclear!

[ttktjmt]

Thanks! Just sent a message.