kscale 0.0.1__tar.gz → 0.0.3__tar.gz

kscale-0.0.3/PKG-INFO

@@ -0,0 +1,52 @@
+Metadata-Version: 2.1
+Name: kscale
+Version: 0.0.3
+Summary: The kscale project
+Home-page: https://github.com/kscalelabs/kscale
+Author: Benjamin Bolte
+Requires-Python: >=3.11
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: omegaconf
+Requires-Dist: httpx
+Requires-Dist: requests
+Provides-Extra: dev
+Requires-Dist: black; extra == "dev"
+Requires-Dist: darglint; extra == "dev"
+Requires-Dist: mypy; extra == "dev"
+Requires-Dist: pytest; extra == "dev"
+Requires-Dist: ruff; extra == "dev"
+Requires-Dist: datamodel-code-generator; extra == "dev"
+
+<p align="center">
+  <picture>
+    <img alt="K-Scale Open Source Robotics" src="https://media.kscale.dev/kscale-open-source-header.png" style="max-width: 100%;">
+  </picture>
+</p>
+
+<div align="center">
+
+[![License](https://img.shields.io/badge/license-MIT-green)](https://github.com/kscalelabs/ksim/blob/main/LICENSE)
+[![Discord](https://img.shields.io/discord/1224056091017478166)](https://discord.gg/k5mSvCkYQh)
+[![Wiki](https://img.shields.io/badge/wiki-humanoids-black)](https://humanoids.wiki)
+<br />
+[![python](https://img.shields.io/badge/-Python_3.11-blue?logo=python&logoColor=white)](https://github.com/pre-commit/pre-commit)
+[![black](https://img.shields.io/badge/Code%20Style-Black-black.svg?labelColor=gray)](https://black.readthedocs.io/en/stable/)
+[![ruff](https://img.shields.io/badge/Linter-Ruff-red.svg?labelColor=gray)](https://github.com/charliermarsh/ruff)
+<br />
+[![Python Checks](https://github.com/kscalelabs/kscale/actions/workflows/test.yml/badge.svg)](https://github.com/kscalelabs/kscale/actions/workflows/test.yml)
+[![Publish Python Package](https://github.com/kscalelabs/kscale/actions/workflows/publish.yml/badge.svg)](https://github.com/kscalelabs/kscale/actions/workflows/publish.yml)
+
+</div>
+
+# K-Scale Command Line Interface
+
+This is a command line tool for interacting with various services provided by K-Scale Labs, such as:
+
+- [K-Scale Store](https://kscale.store/)
+
+## Installation
+
+```bash
+pip install kscale
+```

kscale-0.0.3/README.md

@@ -0,0 +1,32 @@
+<p align="center">
+  <picture>
+    <img alt="K-Scale Open Source Robotics" src="https://media.kscale.dev/kscale-open-source-header.png" style="max-width: 100%;">
+  </picture>
+</p>
+
+<div align="center">
+
+[![License](https://img.shields.io/badge/license-MIT-green)](https://github.com/kscalelabs/ksim/blob/main/LICENSE)
+[![Discord](https://img.shields.io/discord/1224056091017478166)](https://discord.gg/k5mSvCkYQh)
+[![Wiki](https://img.shields.io/badge/wiki-humanoids-black)](https://humanoids.wiki)
+<br />
+[![python](https://img.shields.io/badge/-Python_3.11-blue?logo=python&logoColor=white)](https://github.com/pre-commit/pre-commit)
+[![black](https://img.shields.io/badge/Code%20Style-Black-black.svg?labelColor=gray)](https://black.readthedocs.io/en/stable/)
+[![ruff](https://img.shields.io/badge/Linter-Ruff-red.svg?labelColor=gray)](https://github.com/charliermarsh/ruff)
+<br />
+[![Python Checks](https://github.com/kscalelabs/kscale/actions/workflows/test.yml/badge.svg)](https://github.com/kscalelabs/kscale/actions/workflows/test.yml)
+[![Publish Python Package](https://github.com/kscalelabs/kscale/actions/workflows/publish.yml/badge.svg)](https://github.com/kscalelabs/kscale/actions/workflows/publish.yml)
+
+</div>
+
+# K-Scale Command Line Interface
+
+This is a command line tool for interacting with various services provided by K-Scale Labs, such as:
+
+- [K-Scale Store](https://kscale.store/)
+
+## Installation
+
+```bash
+pip install kscale
+```

kscale-0.0.3/kscale/__init__.py

@@ -0,0 +1 @@
+__version__ = "0.0.3"

{kscale-0.0.1 → kscale-0.0.3}/kscale/conf.py

@@ -17,21 +17,13 @@ def get_path() -> Path:
 
 @dataclass
 class StoreSettings:
-    api_key: str = field(default=II("oc.env:KSCALE_API_KEY"))
-
-    def get_api_key(self) -> str:
-        try:
-            return self.api_key
-        except AttributeError:
-            raise ValueError(
-                "API key not found! Get one here and set it as the `KSCALE_API_KEY` "
-                "environment variable: https://kscale.store/keys"
-            )
+    api_key: str = field(default=II("oc.env:KSCALE_API_KEY,"))
+    cache_dir: str = field(default=II("oc.env:KSCALE_CACHE_DIR,'~/.kscale/cache/'"))
 
 
 @dataclass
 class Settings:
-    store: StoreSettings = StoreSettings()
+    store: StoreSettings = field(default_factory=StoreSettings)
 
     def save(self) -> None:
         (dir_path := get_path()).mkdir(parents=True, exist_ok=True)

kscale-0.0.3/kscale/formats/mjcf.py

@@ -0,0 +1,509 @@
+"""Defines common types and functions for exporting MJCF files.
+
+API reference:
+https://github.com/google-deepmind/mujoco/blob/main/src/xml/xml_native_writer.cc#L780
+
+Todo:
+    1. Clean up the inertia config
+    2. Add visual and imu support
+"""
+
+import glob
+import os
+import shutil
+import xml.etree.ElementTree as ET
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Literal
+
+
+@dataclass
+class Compiler:
+    coordinate: Literal["local", "global"] | None = None
+    angle: Literal["radian", "degree"] = "radian"
+    meshdir: str | None = None
+    eulerseq: Literal["xyz", "zxy", "zyx", "yxz", "yzx", "xzy"] | None = None
+    autolimits: bool | None = None
+
+    def to_xml(self, compiler: ET.Element | None = None) -> ET.Element:
+        if compiler is None:
+            compiler = ET.Element("compiler")
+        if self.coordinate is not None:
+            compiler.set("coordinate", self.coordinate)
+        compiler.set("angle", self.angle)
+        if self.meshdir is not None:
+            compiler.set("meshdir", self.meshdir)
+        if self.eulerseq is not None:
+            compiler.set("eulerseq", self.eulerseq)
+        if self.autolimits is not None:
+            compiler.set("autolimits", str(self.autolimits).lower())
+        return compiler
+
+
+@dataclass
+class Mesh:
+    name: str
+    file: str
+    scale: tuple[float, float, float] | None = None
+
+    def to_xml(self, root: ET.Element | None = None) -> ET.Element:
+        mesh = ET.Element("mesh") if root is None else ET.SubElement(root, "mesh")
+        mesh.set("name", self.name)
+        mesh.set("file", self.file)
+        if self.scale is not None:
+            mesh.set("scale", " ".join(map(str, self.scale)))
+        return mesh
+
+
+@dataclass
+class Joint:
+    name: str | None = None
+    type: Literal["hinge", "slide", "ball", "free"] | None = None
+    pos: tuple[float, float, float] | None = None
+    axis: tuple[float, float, float] | None = None
+    limited: bool | None = None
+    range: tuple[float, float] | None = None
+    damping: float | None = None
+    stiffness: float | None = None
+    armature: float | None = None
+    frictionloss: float | None = None
+
+    def to_xml(self, root: ET.Element | None = None) -> ET.Element:
+        joint = ET.Element("joint") if root is None else ET.SubElement(root, "joint")
+        if self.name is not None:
+            joint.set("name", self.name)
+        if self.type is not None:
+            joint.set("type", self.type)
+        if self.pos is not None:
+            joint.set("pos", " ".join(map(str, self.pos)))
+        if self.axis is not None:
+            joint.set("axis", " ".join(map(str, self.axis)))
+        if self.range is not None:
+            joint.set("range", " ".join(map(str, self.range)))
+        if self.limited is not None:
+            joint.set("limited", str(self.limited).lower())
+        if self.damping is not None:
+            joint.set("damping", str(self.damping))
+        if self.stiffness is not None:
+            joint.set("stiffness", str(self.stiffness))
+        if self.armature is not None:
+            joint.set("armature", str(self.armature))
+        if self.frictionloss is not None:
+            joint.set("frictionloss", str(self.frictionloss))
+        return joint
+
+
+@dataclass
+class Geom:
+    name: str | None = None
+    type: Literal["plane", "sphere", "cylinder", "box", "capsule", "ellipsoid", "mesh"] | None = None
+    plane: str | None = None
+    rgba: tuple[float, float, float, float] | None = None
+    pos: tuple[float, float, float] | None = None
+    quat: tuple[float, float, float, float] | None = None
+    matplane: str | None = None
+    material: str | None = None
+    condim: int | None = None
+    contype: int | None = None
+    conaffinity: int | None = None
+    size: tuple[float, float, float] | None = None
+    friction: tuple[float, float, float] | None = None
+    solref: tuple[float, float] | None = None
+    density: float | None = None
+
+    def to_xml(self, root: ET.Element | None = None) -> ET.Element:
+        geom = ET.Element("geom") if root is None else ET.SubElement(root, "geom")
+        if self.name is not None:
+            geom.set("name", self.name)
+        if self.type is not None:
+            geom.set("type", self.type)
+        if self.rgba is not None:
+            geom.set("rgba", " ".join(map(str, self.rgba)))
+        if self.pos is not None:
+            geom.set("pos", " ".join(map(str, self.pos)))
+        if self.quat is not None:
+            geom.set("quat", " ".join(map(str, self.quat)))
+        if self.matplane is not None:
+            geom.set("matplane", self.matplane)
+        if self.material is not None:
+            geom.set("material", self.material)
+        if self.condim is not None:
+            geom.set("condim", str(self.condim))
+        if self.contype is not None:
+            geom.set("contype", str(self.contype))
+        if self.conaffinity is not None:
+            geom.set("conaffinity", str(self.conaffinity))
+        if self.plane is not None:
+            geom.set("plane", self.plane)
+        if self.size is not None:
+            geom.set("size", " ".join(map(str, self.size)))
+        if self.friction is not None:
+            geom.set("friction", " ".join(map(str, self.friction)))
+        if self.solref is not None:
+            geom.set("solref", " ".join(map(str, self.solref)))
+        if self.density is not None:
+            geom.set("density", str(self.density))
+        return geom
+
+
+@dataclass
+class Body:
+    name: str
+    pos: tuple[float, float, float] | None = field(default=None)
+    quat: tuple[float, float, float, float] | None = field(default=None)
+    geom: Geom | None = field(default=None)
+    joint: Joint | None = field(default=None)
+
+    # TODO - Fix inertia, until then rely on Mujoco's engine
+    # inertial: Inertial = None
+
+    def to_xml(self, root: ET.Element | None = None) -> ET.Element:
+        body = ET.Element("body") if root is None else ET.SubElement(root, "body")
+        body.set("name", self.name)
+        if self.pos is not None:
+            body.set("pos", " ".join(map(str, self.pos)))
+        if self.quat is not None:
+            body.set("quat", " ".join(f"{x:.8g}" for x in self.quat))
+        if self.joint is not None:
+            self.joint.to_xml(body)
+        if self.geom is not None:
+            self.geom.to_xml(body)
+        return body
+
+
+@dataclass
+class Flag:
+    frictionloss: str | None = None
+    # removed at 3.1.4
+    # sensornoise: str | None = None
+
+    def to_xml(self, root: ET.Element | None = None) -> ET.Element:
+        flag = ET.Element("flag") if root is None else ET.SubElement(root, "flag")
+        if self.frictionloss is not None:
+            flag.set("frictionloss", self.frictionloss)
+        return flag
+
+
+@dataclass
+class Option:
+    timestep: float | None = None
+    viscosity: float | None = None
+    iterations: int | None = None
+    solver: Literal["PGS", "CG", "Newton"] | None = None
+    gravity: tuple[float, float, float] | None = None
+    flag: Flag | None = None
+
+    def to_xml(self, root: ET.Element | None = None) -> ET.Element:
+        if root is None:
+            option = ET.Element("option")
+        else:
+            option = ET.SubElement(root, "option")
+        if self.iterations is not None:
+            option.set("iterations", str(self.iterations))
+        if self.timestep is not None:
+            option.set("timestep", str(self.timestep))
+        if self.viscosity is not None:
+            option.set("viscosity", str(self.viscosity))
+        if self.solver is not None:
+            option.set("solver", self.solver)
+        if self.gravity is not None:
+            option.set("gravity", " ".join(map(str, self.gravity)))
+        if self.flag is not None:
+            self.flag.to_xml(option)
+        return option
+
+
+@dataclass
+class Motor:
+    name: str | None = None
+    joint: str | None = None
+    ctrlrange: tuple[float, float] | None = None
+    ctrllimited: bool | None = None
+    gear: float | None = None
+
+    def to_xml(self, root: ET.Element | None = None) -> ET.Element:
+        if root is None:
+            motor = ET.Element("motor")
+        else:
+            motor = ET.SubElement(root, "motor")
+        if self.name is not None:
+            motor.set("name", self.name)
+        if self.joint is not None:
+            motor.set("joint", self.joint)
+        if self.ctrllimited is not None:
+            motor.set("ctrllimited", str(self.ctrllimited).lower())
+        if self.ctrlrange is not None:
+            motor.set("ctrlrange", " ".join(map(str, self.ctrlrange)))
+        if self.gear is not None:
+            motor.set("gear", str(self.gear))
+        return motor
+
+
+@dataclass
+class Light:
+    directional: bool = True
+    diffuse: tuple[float, float, float] | None = None
+    specular: tuple[float, float, float] | None = None
+    pos: tuple[float, float, float] | None = None
+    dir: tuple[float, float, float] | None = None
+    castshadow: bool | None = None
+
+    def to_xml(self, root: ET.Element | None = None) -> ET.Element:
+        if root is None:
+            light = ET.Element("light")
+        else:
+            light = ET.SubElement(root, "light")
+        if self.directional is not None:
+            light.set("directional", str(self.directional).lower())
+        if self.diffuse is not None:
+            light.set("diffuse", " ".join(map(str, self.diffuse)))
+        if self.specular is not None:
+            light.set("specular", " ".join(map(str, self.specular)))
+        if self.pos is not None:
+            light.set("pos", " ".join(map(str, self.pos)))
+        if self.dir is not None:
+            light.set("dir", " ".join(map(str, self.dir)))
+        if self.castshadow is not None:
+            light.set("castshadow", str(self.castshadow).lower())
+        return light
+
+
+@dataclass
+class Equality:
+    solref: tuple[float, float]
+
+    def to_xml(self, root: ET.Element | None = None) -> ET.Element:
+        equality = ET.Element("equality") if root is None else ET.SubElement(root, "equality")
+        equality.set("solref", " ".join(map(str, self.solref)))
+        return equality
+
+
+@dataclass
+class Site:
+    name: str | None = None
+    size: float | None = None
+    pos: tuple[float, float, float] | None = None
+
+    def to_xml(self, root: ET.Element | None = None) -> ET.Element:
+        site = ET.Element("site") if root is None else ET.SubElement(root, "site")
+        if self.name is not None:
+            site.set("name", self.name)
+        if self.size is not None:
+            site.set("size", str(self.size))
+        if self.pos is not None:
+            site.set("pos", " ".join(map(str, self.pos)))
+        return site
+
+
+@dataclass
+class Default:
+    joint: Joint | None = None
+    geom: Geom | None = None
+    class_: str | None = None
+    motor: Motor | None = None
+    equality: Equality | None = None
+    visual_geom: ET.Element | None = None
+
+    def to_xml(self, default: ET.Element | None = None) -> ET.Element:
+        default = ET.Element("default") if default is None else ET.SubElement(default, "default")
+        if self.joint is not None:
+            self.joint.to_xml(default)
+        if self.geom is not None:
+            self.geom.to_xml(default)
+        if self.class_ is not None:
+            default.set("class", self.class_)
+        if self.motor is not None:
+            self.motor.to_xml(default)
+        if self.equality is not None:
+            self.equality.to_xml(default)
+        if self.visual_geom is not None:
+            default.append(self.visual_geom)
+        return default
+
+
+@dataclass
+class Actuator:
+    motors: list[Motor]
+
+    def to_xml(self, root: ET.Element | None = None) -> ET.Element:
+        actuator = ET.Element("actuator") if root is None else ET.SubElement(root, "actuator")
+        for motor in self.motors:
+            motor.to_xml(actuator)
+        return actuator
+
+
+@dataclass
+class Actuatorpos:
+    name: str | None = None
+    actuator: str | None = None
+    user: str | None = None
+
+    def to_xml(self, root: ET.Element | None = None) -> ET.Element:
+        actuatorpos = ET.Element("actuatorpos") if root is None else ET.SubElement(root, "actuatorpos")
+        if self.name is not None:
+            actuatorpos.set("name", self.name)
+        if self.actuator is not None:
+            actuatorpos.set("actuator", self.actuator)
+        if self.user is not None:
+            actuatorpos.set("user", self.user)
+        return actuatorpos
+
+
+@dataclass
+class Actuatorvel:
+    name: str | None = None
+    actuator: str | None = None
+    user: str | None = None
+
+    def to_xml(self, root: ET.Element | None = None) -> ET.Element:
+        actuatorvel = ET.Element("actuatorvel") if root is None else ET.SubElement(root, "actuatorvel")
+        if self.name is not None:
+            actuatorvel.set("name", self.name)
+        if self.actuator is not None:
+            actuatorvel.set("actuator", self.actuator)
+        if self.user is not None:
+            actuatorvel.set("user", self.user)
+        return actuatorvel
+
+
+@dataclass
+class Actuatorfrc:
+    name: str | None = None
+    actuator: str | None = None
+    user: str | None = None
+    noise: float | None = None
+
+    def to_xml(self, root: ET.Element | None = None) -> ET.Element:
+        actuatorfrc = ET.Element("actuatorfrc") if root is None else ET.SubElement(root, "actuatorfrc")
+        if self.name is not None:
+            actuatorfrc.set("name", self.name)
+        if self.actuator is not None:
+            actuatorfrc.set("actuator", self.actuator)
+        if self.user is not None:
+            actuatorfrc.set("user", self.user)
+        if self.noise is not None:
+            actuatorfrc.set("noise", str(self.noise))
+        return actuatorfrc
+
+
+@dataclass
+class Sensor:
+    actuatorpos: list[Actuatorpos] | None = None
+    actuatorvel: list[Actuatorvel] | None = None
+    actuatorfrc: list[Actuatorfrc] | None = None
+
+    def to_xml(self, root: ET.Element | None = None) -> ET.Element:
+        sensor = ET.Element("sensor") if root is None else ET.SubElement(root, "sensor")
+        if self.actuatorpos is not None:
+            for actuatorpos in self.actuatorpos:
+                actuatorpos.to_xml(sensor)
+        if self.actuatorvel is not None:
+            for actuatorvel in self.actuatorvel:
+                actuatorvel.to_xml(sensor)
+        if self.actuatorfrc is not None:
+            for actuatorfrc in self.actuatorfrc:
+                actuatorfrc.to_xml(sensor)
+        return sensor
+
+
+def _copy_stl_files(source_directory: str | Path, destination_directory: str | Path) -> None:
+    # Ensure the destination directory exists, create if not
+    os.makedirs(destination_directory, exist_ok=True)
+
+    # Use glob to find all .stl files in the source directory
+    pattern = os.path.join(source_directory, "*.stl")
+    for file_path in glob.glob(pattern):
+        destination_path = os.path.join(destination_directory, os.path.basename(file_path))
+        shutil.copy(file_path, destination_path)
+        print(f"Copied {file_path} to {destination_path}")
+
+
+def _remove_stl_files(source_directory: str | Path) -> None:
+    for filename in os.listdir(source_directory):
+        if filename.endswith(".stl"):
+            file_path = os.path.join(source_directory, filename)
+            os.remove(file_path)
+
+
+class Robot:
+    """A class to adapt the world in a Mujoco XML file."""
+
+    def __init__(
+        self,
+        robot_name: str,
+        output_dir: str | Path,
+        compiler: Compiler | None = None,
+    ) -> None:
+        """Initialize the robot.
+
+        Args:
+            robot_name (str): The name of the robot.
+            output_dir (str | Path): The output directory.
+            compiler (Compiler, optional): The compiler settings.
+        """
+        self.robot_name = robot_name
+        self.output_dir = Path(output_dir)
+        self.compiler = compiler
+        self._set_clean_up()
+        self.tree = ET.parse(self.output_dir / f"{self.robot_name}.xml")
+
+    def _set_clean_up(self) -> None:
+        try:
+            import mujoco  # type: ignore[import-not-found]
+        except ImportError as e:
+            raise ImportError(
+                "Please install the package with Mujoco as a dependency, using "
+                "`pip install kscale-onshape-library[mujoco]`"
+            ) from e
+
+        # HACK
+        # mujoco has a hard time reading meshes
+        _copy_stl_files(self.output_dir / "meshes", self.output_dir)
+
+        urdf_tree = ET.parse(self.output_dir / f"{self.robot_name}.urdf")
+        root = urdf_tree.getroot()
+
+        tree = ET.ElementTree(root)
+        tree.write(self.output_dir / f"{self.robot_name}.urdf", encoding="utf-8")
+        model = mujoco.MjModel.from_xml_path((self.output_dir / f"{self.robot_name}.urdf").as_posix())
+        mujoco.mj_saveLastXML((self.output_dir / f"{self.robot_name}.xml").as_posix(), model)
+
+        # Removes all the existing files.
+        _remove_stl_files(self.output_dir)
+
+    def adapt_world(self) -> None:
+        root = self.tree.getroot()
+
+        # Turn off internal collisions
+        for element in root:
+            if element.tag == "geom":
+                element.attrib["contype"] = str(1)
+                element.attrib["conaffinity"] = str(0)
+
+        compiler = root.find("compiler")
+        if self.compiler is not None:
+            compiler = self.compiler.to_xml(compiler)
+
+        worldbody = root.find("worldbody")
+        if worldbody is None:
+            raise ValueError("No worldbody found in the XML file")
+
+        new_root_body = Body(name="root", pos=(0, 0, 0), quat=(1, 0, 0, 0)).to_xml()
+
+        # List to store items to be moved to the new root body
+        items_to_move = []
+        # Gather all children (geoms and bodies) that need to be moved under the new root body
+        for element in worldbody:
+            items_to_move.append(element)
+
+        # Move gathered elements to the new root body
+        for item in items_to_move:
+            worldbody.remove(item)
+            new_root_body.append(item)
+
+        # Add the new root body to the worldbody
+        worldbody.append(new_root_body)
+        self.tree = ET.ElementTree(root)
+
+    def save(self, path: str | Path) -> None:
+        self.tree.write(path, encoding="utf-8")

kscale-0.0.3/kscale/requirements.txt

@@ -0,0 +1,8 @@
+# requirements.txt
+
+# Configuration
+omegaconf
+
+# HTTP requests
+httpx
+requests

kscale-0.0.3/kscale/store/bullet/MANIFEST.in

@@ -0,0 +1 @@
+include . *.obj *.urdf

kscale-0.0.3/kscale/store/cli.py

@@ -0,0 +1,35 @@
+"""Defines the top-level KOL CLI."""
+
+import argparse
+import asyncio
+from typing import Sequence
+
+from kscale.store import pybullet, urdf
+
+
+async def main(args: Sequence[str] | None = None) -> None:
+    parser = argparse.ArgumentParser(description="K-Scale OnShape Library", add_help=False)
+    parser.add_argument(
+        "subcommand",
+        choices=[
+            "urdf",
+            "pybullet",
+        ],
+        help="The subcommand to run",
+    )
+    parsed_args, remaining_args = parser.parse_known_args(args)
+
+    match parsed_args.subcommand:
+        case "urdf":
+            await urdf.main(remaining_args)
+        case "pybullet":
+            await pybullet.main(remaining_args)
+
+
+def sync_main(args: Sequence[str] | None = None) -> None:
+    asyncio.run(main(args))
+
+
+if __name__ == "__main__":
+    # python3 -m kscale.store.cli
+    sync_main()

kscale-0.0.3/kscale/store/pybullet.py

@@ -0,0 +1,179 @@
+"""Simple script to interact with a URDF in PyBullet."""
+
+import argparse
+import asyncio
+import itertools
+import logging
+import math
+import time
+from pathlib import Path
+from typing import Sequence
+
+from kscale.store.urdf import download_urdf
+
+logger = logging.getLogger(__name__)
+
+
+async def main(args: Sequence[str] | None = None) -> None:
+    parser = argparse.ArgumentParser(description="Show a URDF in PyBullet")
+    parser.add_argument("listing_id", help="Listing ID for the URDF")
+    parser.add_argument("--dt", type=float, default=0.01, help="Time step")
+    parser.add_argument("-n", "--hide-gui", action="store_true", help="Hide the GUI")
+    parser.add_argument("--no-merge", action="store_true", help="Do not merge fixed links")
+    parser.add_argument("--hide-origin", action="store_true", help="Do not show the origin")
+    parser.add_argument("--show-inertia", action="store_true", help="Visualizes the inertia frames")
+    parser.add_argument("--see-thru", action="store_true", help="Use see-through mode")
+    parser.add_argument("--show-collision", action="store_true", help="Show collision meshes")
+    parsed_args = parser.parse_args(args)
+
+    # Gets the URDF path.
+    urdf_path = await download_urdf(parsed_args.listing_id)
+
+    breakpoint()
+
+    try:
+        import pybullet as p  # type: ignore[import-not-found]
+    except ImportError:
+        raise ImportError("pybullet is required to run this script")
+
+    # Connect to PyBullet.
+    p.connect(p.GUI)
+    p.setGravity(0, 0, -9.81)
+    p.setRealTimeSimulation(0)
+
+    # Turn off panels.
+    if parsed_args.hide_gui:
+        p.configureDebugVisualizer(p.COV_ENABLE_GUI, 0)
+    p.configureDebugVisualizer(p.COV_ENABLE_SEGMENTATION_MARK_PREVIEW, 0)
+    p.configureDebugVisualizer(p.COV_ENABLE_DEPTH_BUFFER_PREVIEW, 0)
+    p.configureDebugVisualizer(p.COV_ENABLE_RGB_BUFFER_PREVIEW, 0)
+
+    # Enable mouse picking.
+    p.configureDebugVisualizer(p.COV_ENABLE_MOUSE_PICKING, 1)
+
+    # Loads the floor plane.
+    floor = p.loadURDF(str((Path(__file__).parent / "bullet" / "plane.urdf").resolve()))
+
+    # Load the robot URDF.
+    start_position = [0.0, 0.0, 1.0]
+    start_orientation = p.getQuaternionFromEuler([0.0, 0.0, 0.0])
+    flags = p.URDF_USE_INERTIA_FROM_FILE
+    if not parsed_args.no_merge:
+        flags |= p.URDF_MERGE_FIXED_LINKS
+    robot = p.loadURDF(str(urdf_path), start_position, start_orientation, flags=flags, useFixedBase=0)
+
+    # Display collision meshes as separate object.
+    if parsed_args.show_collision:
+        collision_flags = p.URDF_USE_INERTIA_FROM_FILE | p.URDF_USE_SELF_COLLISION_EXCLUDE_ALL_PARENTS
+        collision = p.loadURDF(str(urdf_path), start_position, start_orientation, flags=collision_flags, useFixedBase=0)
+
+        # Make collision shapes semi-transparent.
+        joint_ids = [i for i in range(p.getNumJoints(collision))] + [-1]
+        for i in joint_ids:
+            p.changeVisualShape(collision, i, rgbaColor=[1, 0, 0, 0.5])
+
+    # Initializes physics parameters.
+    p.changeDynamics(floor, -1, lateralFriction=1, spinningFriction=-1, rollingFriction=-1)
+    p.setPhysicsEngineParameter(fixedTimeStep=parsed_args.dt, maxNumCmdPer1ms=1000)
+
+    # Shows the origin of the robot.
+    if not parsed_args.hide_origin:
+        p.addUserDebugLine([0, 0, 0], [0.1, 0, 0], [1, 0, 0], parentObjectUniqueId=robot, parentLinkIndex=-1)
+        p.addUserDebugLine([0, 0, 0], [0, 0.1, 0], [0, 1, 0], parentObjectUniqueId=robot, parentLinkIndex=-1)
+        p.addUserDebugLine([0, 0, 0], [0, 0, 0.1], [0, 0, 1], parentObjectUniqueId=robot, parentLinkIndex=-1)
+
+    # Make the robot see-through.
+    joint_ids = [i for i in range(p.getNumJoints(robot))] + [-1]
+    if parsed_args.see_thru:
+        for i in joint_ids:
+            p.changeVisualShape(robot, i, rgbaColor=[1, 1, 1, 0.5])
+
+    def draw_box(pt: list[list[float]], color: tuple[float, float, float], obj_id: int, link_id: int) -> None:
+        assert len(pt) == 8
+        assert all(len(p) == 3 for p in pt)
+
+        mapping = [1, 3, 0, 2]
+        for i in range(4):
+            p.addUserDebugLine(pt[i], pt[i + 4], color, 1, parentObjectUniqueId=obj_id, parentLinkIndex=link_id)
+            p.addUserDebugLine(pt[i], pt[mapping[i]], color, 1, parentObjectUniqueId=obj_id, parentLinkIndex=link_id)
+            p.addUserDebugLine(
+                pt[i + 4], pt[mapping[i] + 4], color, 1, parentObjectUniqueId=obj_id, parentLinkIndex=link_id
+            )
+
+    # Shows bounding boxes around each part of the robot representing the inertia frame.
+    if parsed_args.show_inertia:
+        for i in joint_ids:
+            dynamics_info = p.getDynamicsInfo(robot, i)
+            mass = dynamics_info[0]
+            if mass <= 0:
+                continue
+            inertia = dynamics_info[2]
+            ixx = inertia[0]
+            iyy = inertia[1]
+            izz = inertia[2]
+            box_scale_x = 0.5 * math.sqrt(6 * (izz + iyy - ixx) / mass)
+            box_scale_y = 0.5 * math.sqrt(6 * (izz + ixx - iyy) / mass)
+            box_scale_z = 0.5 * math.sqrt(6 * (ixx + iyy - izz) / mass)
+
+            half_extents = [box_scale_x, box_scale_y, box_scale_z]
+            pt = [
+                [x, y, z]
+                for x, y, z in itertools.product(
+                    [-half_extents[0], half_extents[0]],
+                    [-half_extents[1], half_extents[1]],
+                    [-half_extents[2], half_extents[2]],
+                )
+            ]
+            draw_box(pt, (1, 0, 0), robot, i)
+
+    # Show joint controller.
+    joints: dict[str, int] = {}
+    controls: dict[str, float] = {}
+    for i in range(p.getNumJoints(robot)):
+        joint_info = p.getJointInfo(robot, i)
+        name = joint_info[1].decode("utf-8")
+        joint_type = joint_info[2]
+        joints[name] = i
+        if joint_type == p.JOINT_PRISMATIC:
+            joint_min, joint_max = joint_info[8:10]
+            controls[name] = p.addUserDebugParameter(name, joint_min, joint_max, 0.0)
+        elif joint_type == p.JOINT_REVOLUTE:
+            joint_min, joint_max = joint_info[8:10]
+            controls[name] = p.addUserDebugParameter(name, joint_min, joint_max, 0.0)
+
+    # Run the simulation until the user closes the window.
+    last_time = time.time()
+    prev_control_values = {k: 0.0 for k in controls}
+    while p.isConnected():
+        # Reset the simulation if "r" was pressed.
+        keys = p.getKeyboardEvents()
+        if ord("r") in keys and keys[ord("r")] & p.KEY_WAS_TRIGGERED:
+            p.resetBasePositionAndOrientation(robot, start_position, start_orientation)
+            p.setJointMotorControlArray(
+                robot,
+                range(p.getNumJoints(robot)),
+                p.POSITION_CONTROL,
+                targetPositions=[0] * p.getNumJoints(robot),
+            )
+
+        # Set joint positions.
+        for k, v in controls.items():
+            try:
+                target_position = p.readUserDebugParameter(v)
+                if target_position != prev_control_values[k]:
+                    prev_control_values[k] = target_position
+                    p.setJointMotorControl2(robot, joints[k], p.POSITION_CONTROL, target_position)
+            except p.error:
+                logger.debug("Failed to set joint %s", k)
+                pass
+
+        # Step simulation.
+        p.stepSimulation()
+        cur_time = time.time()
+        time.sleep(max(0, parsed_args.dt - (cur_time - last_time)))
+        last_time = cur_time
+
+
+if __name__ == "__main__":
+    # python -m kscale.store.pybullet
+    asyncio.run(main())

kscale-0.0.3/kscale/store/urdf.py

@@ -0,0 +1,213 @@
+"""Utility functions for managing artifacts in the K-Scale store."""
+
+import argparse
+import asyncio
+import logging
+import os
+import shutil
+import sys
+import tarfile
+from pathlib import Path
+from typing import Literal, Sequence, get_args
+
+import httpx
+import requests
+
+from kscale.conf import Settings
+from kscale.store.gen.api import UrdfResponse
+
+# Set up logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+
+def get_api_key() -> str:
+    api_key = Settings.load().store.api_key
+    if not api_key:
+        raise ValueError(
+            "API key not found! Get one here and set it as the `KSCALE_API_KEY` environment variable or in your "
+            "config file: https://kscale.store/keys"
+        )
+    return api_key
+
+
+def get_cache_dir() -> Path:
+    return Path(Settings.load().store.cache_dir).expanduser().resolve()
+
+
+def get_listing_dir(listing_id: str) -> Path:
+    (cache_dir := get_cache_dir() / listing_id).mkdir(parents=True, exist_ok=True)
+    return cache_dir
+
+
+def fetch_urdf_info(listing_id: str) -> UrdfResponse:
+    url = f"https://api.kscale.store/urdf/info/{listing_id}"
+    headers = {
+        "Authorization": f"Bearer {get_api_key()}",
+    }
+    response = requests.get(url, headers=headers)
+    response.raise_for_status()
+    return UrdfResponse(**response.json())
+
+
+async def download_artifact(artifact_url: str, cache_dir: Path) -> Path:
+    filename = os.path.join(cache_dir, artifact_url.split("/")[-1])
+    headers = {
+        "Authorization": f"Bearer {get_api_key()}",
+    }
+
+    if not os.path.exists(filename):
+        logger.info("Downloading artifact from %s", artifact_url)
+
+        async with httpx.AsyncClient() as client:
+            response = await client.get(artifact_url, headers=headers)
+            response.raise_for_status()
+            with open(filename, "wb") as f:
+                for chunk in response.iter_bytes(chunk_size=8192):
+                    f.write(chunk)
+            logger.info("Artifact downloaded to %s", filename)
+    else:
+        logger.info("Artifact already cached at %s", filename)
+
+    # Extract the .tgz file
+    extract_dir = cache_dir / os.path.splitext(os.path.basename(filename))[0]
+    if not extract_dir.exists():
+        logger.info("Extracting %s to %s", filename, extract_dir)
+        with tarfile.open(filename, "r:gz") as tar:
+            tar.extractall(path=extract_dir)
+    else:
+        logger.info("Artifact already extracted at %s", extract_dir)
+
+    return extract_dir
+
+
+def create_tarball(folder_path: str | Path, output_filename: str, cache_dir: Path) -> str:
+    tarball_path = os.path.join(cache_dir, output_filename)
+    with tarfile.open(tarball_path, "w:gz") as tar:
+        for root, _, files in os.walk(folder_path):
+            for file in files:
+                file_path = os.path.join(root, file)
+                arcname = os.path.relpath(file_path, start=folder_path)
+                tar.add(file_path, arcname=arcname)
+                logger.info("Added %s as %s", file_path, arcname)
+    logger.info("Created tarball %s", tarball_path)
+    return tarball_path
+
+
+async def upload_artifact(tarball_path: str, listing_id: str) -> None:
+    url = f"https://api.kscale.store/urdf/upload/{listing_id}"
+    headers = {
+        "Authorization": f"Bearer {get_api_key()}",
+    }
+
+    async with httpx.AsyncClient() as client:
+        with open(tarball_path, "rb") as f:
+            files = {"file": (f.name, f, "application/gzip")}
+            response = await client.post(url, headers=headers, files=files)
+
+            response.raise_for_status()
+
+    logger.info("Uploaded artifact to %s", url)
+
+
+async def download_urdf(listing_id: str) -> Path:
+    try:
+        urdf_info = fetch_urdf_info(listing_id)
+
+        if urdf_info.urdf is None:
+            breakpoint()
+            raise ValueError(f"No URDF found for listing {listing_id}")
+
+        artifact_url = urdf_info.urdf.url
+        return await download_artifact(artifact_url, get_listing_dir(listing_id))
+
+    except requests.RequestException:
+        logger.exception("Failed to fetch URDF info")
+        raise
+
+
+async def show_urdf_info(listing_id: str) -> None:
+    try:
+        urdf_info = fetch_urdf_info(listing_id)
+
+        if urdf_info.urdf:
+            logger.info("URDF Artifact ID: %s", urdf_info.urdf.artifact_id)
+            logger.info("URDF URL: %s", urdf_info.urdf.url)
+        else:
+            logger.info("No URDF found for listing %s", listing_id)
+    except requests.RequestException:
+        logger.exception("Failed to fetch URDF info")
+        raise
+
+
+async def upload_urdf(listing_id: str, args: Sequence[str] | None = None) -> None:
+    parser = argparse.ArgumentParser(description="Upload a URDF artifact to the K-Scale store")
+    parser.add_argument("folder_path", help="The path to the folder containing the URDF files")
+    parsed_args = parser.parse_args(args)
+    folder_path = Path(parsed_args.folder_path).expanduser().resolve()
+
+    output_filename = f"{listing_id}.tgz"
+    tarball_path = create_tarball(folder_path, output_filename, get_listing_dir(listing_id))
+
+    try:
+        fetch_urdf_info(listing_id)
+        await upload_artifact(tarball_path, listing_id)
+    except requests.RequestException:
+        logger.exception("Failed to upload artifact")
+        raise
+
+
+async def remove_local_urdf(listing_id: str) -> None:
+    try:
+        if listing_id.lower() == "all":
+            cache_dir = get_cache_dir()
+            if cache_dir.exists():
+                logger.info("Removing all local caches at %s", cache_dir)
+                shutil.rmtree(cache_dir)
+            else:
+                logger.error("No local caches found")
+        else:
+            listing_dir = get_listing_dir(listing_id)
+            if listing_dir.exists():
+                logger.info("Removing local cache at %s", listing_dir)
+                shutil.rmtree(listing_dir)
+            else:
+                logger.error("No local cache found for listing %s", listing_id)
+
+    except Exception:
+        logger.error("Failed to remove local cache")
+        raise
+
+
+Command = Literal["download", "info", "upload", "remove-local"]
+
+
+async def main(args: Sequence[str] | None = None) -> None:
+    parser = argparse.ArgumentParser(description="K-Scale URDF Store", add_help=False)
+    parser.add_argument("command", choices=get_args(Command), help="The command to run")
+    parser.add_argument("listing_id", help="The listing ID to operate on")
+    parsed_args, remaining_args = parser.parse_known_args(args)
+
+    command: Command = parsed_args.command
+    listing_id: str = parsed_args.listing_id
+
+    match command:
+        case "download":
+            await download_urdf(listing_id)
+
+        case "info":
+            await show_urdf_info(listing_id)
+
+        case "upload":
+            await upload_urdf(listing_id, remaining_args)
+
+        case "remove-local":
+            await remove_local_urdf(listing_id)
+
+        case _:
+            logger.error("Invalid command")
+            sys.exit(1)
+
+
+if __name__ == "__main__":
+    asyncio.run(main())

kscale-0.0.3/kscale.egg-info/PKG-INFO

@@ -0,0 +1,52 @@
+Metadata-Version: 2.1
+Name: kscale
+Version: 0.0.3
+Summary: The kscale project
+Home-page: https://github.com/kscalelabs/kscale
+Author: Benjamin Bolte
+Requires-Python: >=3.11
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: omegaconf
+Requires-Dist: httpx
+Requires-Dist: requests
+Provides-Extra: dev
+Requires-Dist: black; extra == "dev"
+Requires-Dist: darglint; extra == "dev"
+Requires-Dist: mypy; extra == "dev"
+Requires-Dist: pytest; extra == "dev"
+Requires-Dist: ruff; extra == "dev"
+Requires-Dist: datamodel-code-generator; extra == "dev"
+
+<p align="center">
+  <picture>
+    <img alt="K-Scale Open Source Robotics" src="https://media.kscale.dev/kscale-open-source-header.png" style="max-width: 100%;">
+  </picture>
+</p>
+
+<div align="center">
+
+[![License](https://img.shields.io/badge/license-MIT-green)](https://github.com/kscalelabs/ksim/blob/main/LICENSE)
+[![Discord](https://img.shields.io/discord/1224056091017478166)](https://discord.gg/k5mSvCkYQh)
+[![Wiki](https://img.shields.io/badge/wiki-humanoids-black)](https://humanoids.wiki)
+<br />
+[![python](https://img.shields.io/badge/-Python_3.11-blue?logo=python&logoColor=white)](https://github.com/pre-commit/pre-commit)
+[![black](https://img.shields.io/badge/Code%20Style-Black-black.svg?labelColor=gray)](https://black.readthedocs.io/en/stable/)
+[![ruff](https://img.shields.io/badge/Linter-Ruff-red.svg?labelColor=gray)](https://github.com/charliermarsh/ruff)
+<br />
+[![Python Checks](https://github.com/kscalelabs/kscale/actions/workflows/test.yml/badge.svg)](https://github.com/kscalelabs/kscale/actions/workflows/test.yml)
+[![Publish Python Package](https://github.com/kscalelabs/kscale/actions/workflows/publish.yml/badge.svg)](https://github.com/kscalelabs/kscale/actions/workflows/publish.yml)
+
+</div>
+
+# K-Scale Command Line Interface
+
+This is a command line tool for interacting with various services provided by K-Scale Labs, such as:
+
+- [K-Scale Store](https://kscale.store/)
+
+## Installation
+
+```bash
+pip install kscale
+```

{kscale-0.0.1 → kscale-0.0.3}/kscale.egg-info/SOURCES.txt

@@ -12,11 +12,15 @@ kscale/requirements.txt
 kscale.egg-info/PKG-INFO
 kscale.egg-info/SOURCES.txt
 kscale.egg-info/dependency_links.txt
+kscale.egg-info/entry_points.txt
 kscale.egg-info/requires.txt
 kscale.egg-info/top_level.txt
+kscale/formats/mjcf.py
 kscale/store/__init__.py
-kscale/store/auth.py
+kscale/store/cli.py
+kscale/store/pybullet.py
 kscale/store/urdf.py
+kscale/store/bullet/MANIFEST.in
 kscale/store/gen/__init__.py
 kscale/store/gen/api.py
 tests/test_dummy.py

kscale-0.0.3/kscale.egg-info/entry_points.txt

@@ -0,0 +1,2 @@
+[console_scripts]
+kscale = kscale.store.cli:sync_main

{kscale-0.0.1 → kscale-0.0.3}/kscale.egg-info/requires.txt

@@ -1,4 +1,6 @@
 omegaconf
+httpx
+requests
 
 [dev]
 black

{kscale-0.0.1 → kscale-0.0.3}/setup.cfg

@@ -5,6 +5,10 @@ packages = find:
 exclude = 
 	tests
 
+[options.entry_points]
+console_scripts = 
+	kscale = kscale.store.cli:sync_main
+
 [egg_info]
 tag_build = 
 tag_date = 0

{kscale-0.0.1 → kscale-0.0.3}/setup.py

@@ -1,4 +1,4 @@
-# mypy: disable-error-code="import-untyped"
+# mypy: disable-error-code="import-untyped, import-not-found"
 #!/usr/bin/env python
 """Setup script for the project."""
 
@@ -29,6 +29,7 @@ setup(
     version=version,
     description="The kscale project",
     author="Benjamin Bolte",
+    license_files=("LICENSE",),
     url="https://github.com/kscalelabs/kscale",
     long_description=long_description,
     long_description_content_type="text/markdown",

kscale-0.0.1/PKG-INFO

@@ -1,29 +0,0 @@
-Metadata-Version: 2.1
-Name: kscale
-Version: 0.0.1
-Summary: The kscale project
-Home-page: https://github.com/kscalelabs/kscale
-Author: Benjamin Bolte
-Requires-Python: >=3.11
-Description-Content-Type: text/markdown
-License-File: LICENSE
-Requires-Dist: omegaconf
-Provides-Extra: dev
-Requires-Dist: black; extra == "dev"
-Requires-Dist: darglint; extra == "dev"
-Requires-Dist: mypy; extra == "dev"
-Requires-Dist: pytest; extra == "dev"
-Requires-Dist: ruff; extra == "dev"
-Requires-Dist: datamodel-code-generator; extra == "dev"
-
-# K-Scale Command Line Interface
-
-This is a command line tool for interacting with various services provided by K-Scale Labs, such as:
-
-- [K-Scale Store](https://kscale.store/)
-
-## Installation
-
-```bash
-pip install kscale
-```

kscale-0.0.1/README.md

@@ -1,11 +0,0 @@
-# K-Scale Command Line Interface
-
-This is a command line tool for interacting with various services provided by K-Scale Labs, such as:
-
-- [K-Scale Store](https://kscale.store/)
-
-## Installation
-
-```bash
-pip install kscale
-```

kscale-0.0.1/kscale/requirements.txt

@@ -1,3 +0,0 @@
-# requirements.txt
-
-omegaconf

kscale-0.0.1/kscale/store/auth.py

@@ -1,13 +0,0 @@
-"""Defines utility functions for authenticating the K-Scale Store API."""
-
-from kscale.conf import Settings
-
-
-def get_api_key() -> str:
-    try:
-        return Settings.load().store.api_key
-    except AttributeError:
-        raise ValueError(
-            "API key not found! Get one here and set it as the `KSCALE_API_KEY` "
-            "environment variable: https://kscale.store/keys"
-        )

kscale-0.0.1/kscale/store/urdf.py

@@ -1 +0,0 @@
-"""Utility functions for managing artifacts in the K-Scale store."""

kscale-0.0.1/kscale.egg-info/PKG-INFO

@@ -1,29 +0,0 @@
-Metadata-Version: 2.1
-Name: kscale
-Version: 0.0.1
-Summary: The kscale project
-Home-page: https://github.com/kscalelabs/kscale
-Author: Benjamin Bolte
-Requires-Python: >=3.11
-Description-Content-Type: text/markdown
-License-File: LICENSE
-Requires-Dist: omegaconf
-Provides-Extra: dev
-Requires-Dist: black; extra == "dev"
-Requires-Dist: darglint; extra == "dev"
-Requires-Dist: mypy; extra == "dev"
-Requires-Dist: pytest; extra == "dev"
-Requires-Dist: ruff; extra == "dev"
-Requires-Dist: datamodel-code-generator; extra == "dev"
-
-# K-Scale Command Line Interface
-
-This is a command line tool for interacting with various services provided by K-Scale Labs, such as:
-
-- [K-Scale Store](https://kscale.store/)
-
-## Installation
-
-```bash
-pip install kscale
-```