--- name: kscale-ksim-2bc118524ffe description: '- User asks about robot simulation, humanoid locomotion, or RL policy training' --- # K-Scale ksim Skill > *"RL training library for humanoid locomotion and manipulation. Built on MuJoCo and JAX."* ## Trigger Conditions - User asks about robot simulation, humanoid locomotion, or RL policy training - Questions about MuJoCo/MJX, JAX-based physics simulation - Training walking/manipulation policies for humanoid robots - Sim2Real transfer, domain randomization, curriculum learning ## Overview **ksim** is K-Scale Labs' modular reinforcement learning framework for training robot control policies. It provides: 1. **Physics Engines**: MuJoCo (CPU) and MJX (GPU/JAX-native) 2. **Observation System**: Stateless and stateful observations with noise injection 3. **Reward Functions**: Composable, curriculum-scaled reward components 4. **Action Processing**: Latency modeling, actuator dynamics ## Architecture ``` ┌─────────────────────────────────────────────────────────────────────────┐ │ ksim Control Loop │ │ │ │ ┌──────────────┐ step ┌──────────────┐ observe ┌──────────────┐ │ │ │ PhysicsState │────────▶│ Trajectory │──────────▶│ Observation │ │ │ │ (MJX/MJ) │ │ (batched) │ │ (noisy) │ │ │ └──────┬───────┘ └──────┬───────┘ └──────────────┘ │ │ │ │ │ │ │ reset │ reward │ │ ▼ ▼ │ │ ┌──────────────┐ ┌──────────────┐ │ │ │ Engine │ │ RewardState │ │ │ │ (JIT-compiled)│ │ (component) │ │ │ └──────────────┘ └──────────────┘ │ │ │ │ │ │ actuate │ │ ▼ │ │ ┌──────────────┐ │ │ │ Action │◀─────── policy(observation) │ │ │ (latency) │ │ │ └──────────────┘ │ └─────────────────────────────────────────────────────────────────────────┘ ``` ## Module Structure | Module | Purpose | Key Classes | |--------|---------|-------------| | `engine.py` | Physics stepping | `MjxEngine`, `MujocoEngine`, `PhysicsEngine` | | `observation.py` | State extraction | `Observation`, `StatefulObservation` | | `rewards.py` | Reward computation | `Reward`, `StatefulReward` | | `actions.py` | Action processing | `Action`, latency buffers | | `actuators.py` | Motor dynamics | `Actuators`, `PositionActuators` | | `types.py` | Core dataclasses | `PhysicsState`, `Trajectory`, `RewardState` | | `curriculum.py` | Training progression | `Scale`, curriculum schedules | | `terminations.py` | Episode ending | Termination conditions | | `resets.py` | State initialization | Reset distributions | | `randomization.py` | Domain randomization | Parameter perturbations | ## Behavior Type Taxonomy ### Tree-Sitter Lifted Type Hierarchy ``` RL_BEHAVIOR_INTERFACE (Root) ├─ SIMULATION_BEHAVIOR │ ├─ PHYSICS_STEP: dt × Action → State │ ├─ RESET: InitConfig → State │ ├─ RENDER: State → Visualization │ └─ BATCH_VECTORIZATION: 1 Scene → N Scenes (parallel) │ ├─ POLICY_BEHAVIOR │ ├─ ACTOR: (Obs, LSTM_Carry) → (Distribution, LSTM_Carry) │ ├─ CRITIC: (Obs, LSTM_Carry) → (Value, LSTM_Carry) │ └─ ACTION_SAMPLING: Distribution → Action (JAX random) │ ├─ REWARD_BEHAVIOR │ ├─ PENALTY_COMPOSITION: [Penalty] → Scalar │ ├─ JOINT_DEVIATION: Physics → Scalar │ └─ POSTURE_CONSTRAINT: Physics → Scalar │ ├─ TRAINING_BEHAVIOR │ ├─ TRAJECTORY_COLLECTION: Action × Physics → Experience │ ├─ GRADIENT_COMPUTATION: Trajectory → Gradients (PPO) │ ├─ MODEL_UPDATE: Gradients → Model′ │ └─ CHECKPOINT: Model → Disk │ ├─ CONFIGURATION_BEHAVIOR │ ├─ HYPERPARAMETER_SPEC: Type-safe declarative │ └─ ENV_FACTORY: Config → (Model, Physics, Task) │ └─ RECURRENCE_BEHAVIOR ├─ LSTM_STATE_CARRY: Array → LSTM_Carry ├─ STATE_RESET: () → LSTM_Carry └─ STATE_EVOLUTION: (Carry, Obs) → Carry′ ``` ### Type Signature Contracts | Behavior | Input Type | Output Type | JAX/Equinox Traits | |----------|-----------|------------|-------------------| | `Actor.forward` | `(Array[B,O], Array[H])` | `(Distribution, Array[H])` | PyTree, JIT-compiled | | `Critic.forward` | `(Array[B,O], Array[H])` | `(Array[B,1], Array[H])` | PyTree, Differentiable | | `step(action)` | `Action: Array[B,A]` | `State: PhysicsModel` | Vectorized, MJX batch | | `get_rewards()` | `PhysicsModel` | `Array[B,1]` | JAX-pure function | | `sample_action()` | `(Model, PhysicsModel)` | `Action: Array[B,A]` | Random keyed, PRNGKey | | `Config.__init__()` | Keyword args | `Config: dataclass` | Immutable, type-checked | ### Stateless Behaviors (Pure Functions) ```python # Observation: PhysicsState → Array class BasePositionObservation(Observation): def observe(self, state: PhysicsState) -> Array: return state.data.qpos[0:3] # Reward: Trajectory → Array class BaseHeightReward(Reward): def get_reward(self, trajectory: Trajectory) -> Array: height = trajectory.qpos[:, 2] return jnp.exp(-((height - self.target) ** 2) / (2 * self.scale ** 2)) ``` ### Stateful Behaviors (With Carry) ```python # StatefulObservation: (PhysicsState, Carry) → (Array, Carry) class DelayedJointPositionObservation(StatefulObservation): def observe_stateful(self, state, carry): # Ring buffer for action latency simulation new_carry = jnp.roll(carry, 1, axis=0) new_carry = new_carry.at[0].set(state.data.qpos[7:]) return carry[-1], new_carry # StatefulReward: (Trajectory, Carry) → (Array, Carry) class FeetAirTimeReward(StatefulReward): def get_reward_stateful(self, trajectory, carry): # Track contact state over time ... ``` ### Neural Network Behavioral Contracts (Equinox) ```python class Model(eqx.Module): actor: Actor # Stochastic Policy Behavior critic: Critic # Value Estimation Behavior class Actor(eqx.Module): """Behavioral Contract: (Obs, LSTM_Carry) → (Distribution, LSTM_Carry)""" def forward(self, obs_n: Array, carry: Array) -> tuple[Distribution, Array]: ... class Critic(eqx.Module): """Behavioral Contract: (Obs, LSTM_Carry) → (Value, LSTM_Carry)""" def forward(self, obs_n: Array, carry: Array) -> tuple[Array, Array]: ... ``` ## Key Patterns ### 1. JIT Compilation with Equinox ```python @eqx.filter_jit def step(self, action, physics_model, physics_state, curriculum_level, rng): # Efficient GPU execution via JAX tracing ... ``` ### 2. Exponential Kernel Rewards ```python def exp_kernel(x, scale): return jnp.exp(-(x ** 2) / (2 * scale ** 2)) ``` ### 3. Curriculum Scaling ```python class Scale: def __call__(self, curriculum_level: Array) -> Array: # Modulate reward/observation based on training progress ... ``` ## GF(3) Trit Assignment ``` Trit: 0 (ERGODIC) Role: Infrastructure/Coordination Color: #25BC3D URI: skill://kscale-ksim#25BC3D ``` ### Balanced Triads ``` kscale-ksim (0) ⊗ kscale-kos (-1) ⊗ gym (+1) = 0 ✓ kscale-ksim (0) ⊗ jax-rl (-1) ⊗ mujoco-playground (+1) = 0 ✓ ``` ## Related Skills - `kscale-kos`: K-Scale Operating System (firmware layer) - `kscale-kinfer`: Model inference engine - `kscale-urdf`: Robot description conversion - `gym`: OpenAI Gym environments - `jax`: JAX numerical computing ## Key Contributors (Cognitive Superposition) | Contributor | Focus Areas | Commits | |------------|-------------|---------| | **codekansas** (Ben Bolte) | Architecture, rewards, training | 1475+ | | **b-vm** | Randomization, disturbances | 500+ | | **WT-MM** (Wesley Maa) | Tooling, visualization | 300+ | | **alik-git** (Ali Kuwajerwala) | Integration, testing | 200+ | ## Commands ```bash # Install ksim pip install ksim # Train a walking policy (RTX 4090: ~30 min for 80 steps) python -m ksim.train --config configs/kbot_walk.yaml # Visualize trained policy python -m ksim.vis --checkpoint path/to/model.ckpt ``` ## References - [kscalelabs/ksim](https://github.com/kscalelabs/ksim) - Main repository - [kscalelabs/ksim-gym](https://github.com/kscalelabs/ksim-gym) - Training harness - [MuJoCo Playground](https://arxiv.org/html/2502.08844v1) - Design influences - [docs.kscale.dev](https://docs.kscale.dev) - Official documentation ## Narya Compatibility (Structure-Aware Diffing) | Field | Definition | |-------|------------| | `before` | `PhysicsState` at timestep t (qpos, qvel, control) | | `after` | `PhysicsState` at timestep t+1 after action execution | | `delta` | `Trajectory` segment: the action taken + reward received | | `birth` | Initial `PhysicsState` from `reset()` with domain randomization | | `impact` | 1 if episode terminated (fall, out-of-bounds), 0 otherwise | ### Behavior Type Diffing ```python @dataclass class KsimNaryaEvent: """Structure-aware diff for ksim state transitions.""" event_id: str before: PhysicsState # State before action after: PhysicsState # State after action delta: TrajectorySegment # Action + reward + info trit: int # GF(3): -1=penalty, 0=neutral, +1=reward @property def impact(self) -> int: """1 if state change is significant (termination/reset).""" return 1 if self.delta.done else 0 def to_jsonl(self) -> str: return json.dumps({ "event_id": self.event_id, "before_hash": hash_state(self.before), "after_hash": hash_state(self.after), "delta": {"action": self.delta.action.tolist(), "reward": float(self.delta.reward)}, "trit": self.trit, "impact": self.impact }) ``` ### Replay Determinism ```python # Same seed → same trajectory (critical for sim2real debugging) def replay_episode(seed: int, policy: Model) -> list[KsimNaryaEvent]: rng = jax.random.PRNGKey(seed) state = env.reset(rng) # birth events = [] for t in range(max_steps): rng, action_rng = jax.random.split(rng) action = policy.sample(state.obs, action_rng) before = state state, reward, done, info = env.step(action) events.append(KsimNaryaEvent( event_id=f"step_{t}", before=before, after=state, delta=TrajectorySegment(action, reward, done, info), trit=sign(reward) # -1, 0, +1 )) if done: break return events ``` ## ACSet Schema ```julia @present SchKsim(FreeSchema) begin PhysicsState::Ob Trajectory::Ob Observation::Ob Reward::Ob Action::Ob step::Hom(Action, PhysicsState) observe::Hom(PhysicsState, Observation) reward::Hom(Trajectory, Reward) StateData::AttrType qpos::Attr(PhysicsState, StateData) qvel::Attr(PhysicsState, StateData) end ``` --- Source attribution: plurigrid Original project: https://github.com/plurigrid/asi/blob/293ce9f76e3abcd5ba18db284688bc0d9ac6691d/skills/kscale-ksim/SKILL.md License: MIT License information: https://opensource.org/license/mit Original upstream instructions; downloaded from Mega AI Hub. ## Upstream license notice MIT License Copyright (c) 2025 SkillCreator.ai Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions: The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software. THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. # Mega AI Hub resource metadata Creator / publisher: plurigrid Original project: https://github.com/plurigrid/asi/blob/293ce9f76e3abcd5ba18db284688bc0d9ac6691d/skills/kscale-ksim/SKILL.md Attribution basis: Source repository record Fields of work: Data & Analytics, Development & Engineering Roles: Backend Developer, Data Scientist, ML / AI Engineer Specialties: Backend, Design, Word Documents, Frontend, Git, Testing Tools: Jax, Mujoco Country applicability: Country not specified Classification basis: Source catalog labels Instruction SHA-256: e3c0cb4b944249f4bb4349e92cff6f28a0eac1ea14ccda7cc33950c8f956b418