kinfer 0.4.0__cp311-cp311-macosx_11_0_arm64.whl → 0.4.1__cp311-cp311-macosx_11_0_arm64.whl

kinfer/common/types.py

@@ -9,3 +9,4 @@ from pydantic import BaseModel
 
 class Metadata(BaseModel):
     joint_names: list[str]
+    num_commands: int | None

kinfer/export/common.py

@@ -4,6 +4,7 @@
 def get_shape(
     name: str,
     num_joints: int | None = None,
+    num_commands: int | None = None,
     carry_shape: tuple[int, ...] | None = None,
 ) -> tuple[int, ...]:
     match name:
@@ -26,6 +27,11 @@ def get_shape(
         case "gyroscope":
             return (3,)
 
+        case "command":
+            if num_commands is None:
+                raise ValueError("`num_commands` must be provided when using `command`")
+            return (num_commands,)
+
         case "carry":
             if carry_shape is None:
                 raise ValueError("`carry_shape` must be provided for `carry`")

kinfer/export/jax.py

@@ -19,6 +19,7 @@ def export_fn(
     model: Wrapped,
     *,
     num_joints: int | None = None,
+    num_commands: int | None = None,
     carry_shape: tuple[int, ...] | None = None,
     opset: int = 13,
 ) -> ModelProto:
@@ -35,6 +36,7 @@ def export_fn(
         shape = get_shape(
             name,
             num_joints=num_joints,
+            num_commands=num_commands,
             carry_shape=carry_shape,
         )
         tf_args.append(tf.TensorSpec(shape, tf.float32, name=name))

kinfer/export/pytorch.py

@@ -19,6 +19,7 @@ def export_fn(
     model: torch.jit.ScriptFunction,
     *,
     num_joints: int | None = None,
+    num_commands: int | None = None,
     carry_shape: tuple[int, ...] | None = None,
 ) -> ModelProto:
     """Exports a PyTorch function to ONNX.
@@ -26,6 +27,7 @@ def export_fn(
     Args:
         model: The model to export.
         num_joints: The number of joints in the model.
+        num_commands: The number of commands in the model.
         carry_shape: The shape of the carry tensor.
 
     Returns:
@@ -43,6 +45,7 @@ def export_fn(
         shape = get_shape(
             name,
             num_joints=num_joints,
+            num_commands=num_commands,
             carry_shape=carry_shape,
         )
         args.append(torch.zeros(shape))

kinfer/export/serialize.py

@@ -18,7 +18,8 @@ def pack(
     init_fn: ModelProto,
     step_fn: ModelProto,
     joint_names: list[str],
-    carry_shape: tuple[int, ...],
+    num_commands: int | None = None,
+    carry_shape: tuple[int, ...] | None = None,
 ) -> bytes:
     """Packs the initialization function and step function into a directory.
 
@@ -27,8 +28,8 @@ def pack(
         step_fn: The step function.
         joint_names: The list of joint names, in the order that the model
             expects them to be provided.
+        num_commands: The number of commands in the model.
         carry_shape: The shape of the carry tensor.
-        root_dir: The root directory of the model.
     """
     num_joints = len(joint_names)
 
@@ -37,6 +38,10 @@ def pack(
         raise ValueError(f"`init` function should not have any inputs! Got {len(init_fn.graph.input)}")
     if len(init_fn.graph.output) != 1:
         raise ValueError(f"`init` function should have exactly 1 output! Got {len(init_fn.graph.output)}")
+    init_carry = init_fn.graph.output[0]
+    init_carry_shape = tuple(dim.dim_value for dim in init_carry.type.tensor_type.shape.dim)
+    if carry_shape is not None and init_carry_shape != carry_shape:
+        raise ValueError(f"Expected carry shape {carry_shape} for output `{init_carry.name}`, got {init_carry_shape}")
 
     # Checks the `step` function.
     for step_input in step_fn.graph.input:
@@ -45,6 +50,7 @@ def pack(
         expected_shape = get_shape(
             step_input.name,
             num_joints=num_joints,
+            num_commands=num_commands,
             carry_shape=carry_shape,
         )
         if shape != expected_shape:
@@ -53,19 +59,20 @@ def pack(
     if len(step_fn.graph.output) != 2:
         raise ValueError(f"Step function must have exactly 2 outputs, got {len(step_fn.graph.output)}")
 
-    model_output = step_fn.graph.output[0]
-    output_shape = tuple(dim.dim_value for dim in model_output.type.tensor_type.shape.dim)
-    if output_shape != (num_joints,):
-        raise ValueError(f"Expected output shape {num_joints} for output `{model_output.name}`, got {output_shape}")
+    output_actions = step_fn.graph.output[0]
+    actions_shape = tuple(dim.dim_value for dim in output_actions.type.tensor_type.shape.dim)
+    if actions_shape != (num_joints,):
+        raise ValueError(f"Expected output shape {num_joints} for output `{output_actions.name}`, got {actions_shape}")
 
-    model_carry = step_fn.graph.output[1]
-    output_carry_shape = tuple(dim.dim_value for dim in model_carry.type.tensor_type.shape.dim)
-    if output_carry_shape != carry_shape:
-        raise ValueError(f"Expected carry shape {carry_shape} for output `{model_carry.name}`, got {carry_shape}")
+    output_carry = step_fn.graph.output[1]
+    output_carry_shape = tuple(dim.dim_value for dim in output_carry.type.tensor_type.shape.dim)
+    if output_carry_shape != init_carry_shape:
+        raise ValueError(f"Expected carry shape {init_carry_shape} for output carry, got {output_carry_shape}")
 
     # Builds the metadata object.
     metadata = Metadata(
         joint_names=joint_names,
+        num_commands=num_commands,
     )
 
     buffer = io.BytesIO()

kinfer/rust/Cargo.toml

@@ -1,8 +1,18 @@
 [package]
 
 name = "kinfer"
-version = "0.1.0"
-edition = "2021"
+version.workspace = true
+edition.workspace = true
+description.workspace = true
+authors.workspace = true
+repository.workspace = true
+license.workspace = true
+readme.workspace = true
+
+[lib]
+
+name = "kinfer"
+crate-type = ["cdylib", "rlib"]
 
 [dependencies]
 

kinfer/rust/src/model.rs

@@ -16,6 +16,7 @@ use tokio::io::AsyncReadExt;
 #[derive(Debug, Deserialize)]
 struct ModelMetadata {
     joint_names: Vec<String>,
+    num_commands: Option<usize>,
 }
 
 impl ModelMetadata {
@@ -45,6 +46,7 @@ pub trait ModelProvider: Send + Sync {
     async fn get_projected_gravity(&self) -> Result<Array<f32, IxDyn>, ModelError>;
     async fn get_accelerometer(&self) -> Result<Array<f32, IxDyn>, ModelError>;
     async fn get_gyroscope(&self) -> Result<Array<f32, IxDyn>, ModelError>;
+    async fn get_command(&self) -> Result<Array<f32, IxDyn>, ModelError>;
     async fn get_carry(&self, carry: Array<f32, IxDyn>) -> Result<Array<f32, IxDyn>, ModelError>;
     async fn take_action(
         &self,
@@ -134,7 +136,7 @@ impl ModelRunner {
             .to_vec();
 
         // Validate step_fn inputs and outputs
-        Self::validate_step_fn(&step_session, metadata.joint_names.len(), &carry_shape)?;
+        Self::validate_step_fn(&step_session, &metadata, &carry_shape)?;
 
         Ok(Self {
             init_session,
@@ -146,7 +148,7 @@ impl ModelRunner {
 
     fn validate_step_fn(
         session: &Session,
-        num_joints: usize,
+        metadata: &ModelMetadata,
         carry_shape: &[i64],
     ) -> Result<(), Box<dyn std::error::Error>> {
         // Validate inputs
@@ -158,6 +160,7 @@ impl ModelRunner {
 
             match input.name.as_str() {
                 "joint_angles" | "joint_angular_velocities" => {
+                    let num_joints = metadata.joint_names.len();
                     if *dims != vec![num_joints as i64] {
                         return Err(format!(
                             "Expected shape [{num_joints}] for input `{}`, got {:?}",
@@ -175,6 +178,16 @@ impl ModelRunner {
                         .into());
                     }
                 }
+                "command" => {
+                    let num_commands = metadata.num_commands.ok_or("num_commands is not set")?;
+                    if *dims != vec![num_commands as i64] {
+                        return Err(format!(
+                            "Expected shape [{num_commands}] for input `{}`, got {:?}",
+                            input.name, dims
+                        )
+                        .into());
+                    }
+                }
                 "carry" => {
                     if dims != carry_shape {
                         return Err(format!(
@@ -197,6 +210,7 @@ impl ModelRunner {
             .output_type
             .tensor_dimensions()
             .ok_or("Missing tensor type")?;
+        let num_joints = metadata.joint_names.len();
         if *output_shape != vec![num_joints as i64] {
             return Err(format!(
                 "Expected output shape [{num_joints}], got {:?}",
@@ -253,6 +267,7 @@ impl ModelRunner {
                 "projected_gravity" => futures.push(self.provider.get_projected_gravity()),
                 "accelerometer" => futures.push(self.provider.get_accelerometer()),
                 "gyroscope" => futures.push(self.provider.get_gyroscope()),
+                "command" => futures.push(self.provider.get_command()),
                 "carry" => futures.push(self.provider.get_carry(carry.clone())),
                 _ => return Err(format!("Unknown input name: {}", name).into()),
             }

kinfer/rust_bindings/rust_bindings.pyi

@@ -13,6 +13,7 @@ class ModelProviderABC:
     def get_projected_gravity(self) -> numpy.typing.NDArray[numpy.float32]: ...
     def get_accelerometer(self) -> numpy.typing.NDArray[numpy.float32]: ...
     def get_gyroscope(self) -> numpy.typing.NDArray[numpy.float32]: ...
+    def get_command(self) -> numpy.typing.NDArray[numpy.float32]: ...
     def take_action(self, joint_names:typing.Sequence[builtins.str], action:numpy.typing.NDArray[numpy.float32]) -> None: ...
 
 class PyModelProvider:

kinfer/rust_bindings/src/lib.rs

@@ -69,6 +69,10 @@ impl ModelProviderABC {
         ))
     }
 
+    fn get_command<'py>(&self) -> PyResult<Bound<'py, PyArray1<f32>>> {
+        Err(PyNotImplementedError::new_err("Must override get_command"))
+    }
+
     fn take_action<'py>(
         &self,
         joint_names: Vec<String>,
@@ -166,6 +170,18 @@ impl ModelProvider for PyModelProvider {
         Ok(args)
     }
 
+    async fn get_command(&self) -> Result<Array<f32, IxDyn>, ModelError> {
+        let args = Python::with_gil(|py| -> PyResult<Array<f32, IxDyn>> {
+            let obj = self.obj.clone();
+            let args = ();
+            let result = obj.call_method(py, "get_command", args, None)?;
+            let array = result.extract::<Vec<f32>>(py)?;
+            Ok(Array::from_vec(array).into_dyn())
+        })
+        .map_err(|e| ModelError::Provider(e.to_string()))?;
+        Ok(args)
+    }
+
     async fn get_carry(&self, carry: Array<f32, IxDyn>) -> Result<Array<f32, IxDyn>, ModelError> {
         Ok(carry)
     }

kinfer/rust_bindings.pyi

@@ -13,6 +13,7 @@ class ModelProviderABC:
     def get_projected_gravity(self) -> numpy.typing.NDArray[numpy.float32]: ...
     def get_accelerometer(self) -> numpy.typing.NDArray[numpy.float32]: ...
     def get_gyroscope(self) -> numpy.typing.NDArray[numpy.float32]: ...
+    def get_command(self) -> numpy.typing.NDArray[numpy.float32]: ...
     def take_action(self, joint_names:typing.Sequence[builtins.str], action:numpy.typing.NDArray[numpy.float32]) -> None: ...
 
 class PyModelProvider:

{kinfer-0.4.0.dist-info → kinfer-0.4.1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: kinfer
-Version: 0.4.0
+Version: 0.4.1
 Summary: Tool to make it easier to run a model on a real robot
 Home-page: https://github.com/kscalelabs/kinfer.git
 Author: K-Scale Labs

kinfer-0.4.1.dist-info/RECORD

@@ -0,0 +1,26 @@
+kinfer/rust_bindings.cpython-311-darwin.so,sha256=HcN2IJN0pdyOYcReklwbT3ppgCgExjr_pwHY7zayMyw,1933392
+kinfer/requirements.txt,sha256=j08HO4ptA5afuj99j8FlAP2qla5Zf4_OiEBtgAmF7Jg,90
+kinfer/__init__.py,sha256=YbtJIepEE4pbjYbdgFCV-rBP90AnQpBfDaprkflBmEE,99
+kinfer/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+kinfer/rust_bindings.pyi,sha256=9u0_TuHwV9xQH3O5OKbYfUZ_Uj6bsUwEOW0bfZujRMQ,1709
+kinfer/rust/Cargo.toml,sha256=uxjvuM1Sm82UgCdHW3R6PbbiJjXvfE4VCLvjGPqp2mY,606
+kinfer/rust/src/runtime.rs,sha256=eDflFSIa4IxOmG3SInwrHrYc02UzEANZMkK-liA7pVA,3407
+kinfer/rust/src/lib.rs,sha256=Z3dWdhKhqhWqPJee6vWde4ptqquOYW6W9wB0wyaKCyk,71
+kinfer/rust/src/model.rs,sha256=NqxWjSDivkN2mLAGCqytjfYHlx8dpP4zqoatx7pRWg4,11419
+kinfer/common/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+kinfer/common/types.py,sha256=jiuCfoTEumOJm4VZ03VoFVIL5PMJ1tCG7b2Kr66qJQA,176
+kinfer/export/serialize.py,sha256=7rcSrXawVnJinxeJp26wovUwoR04BnJ5CZAV9vwKUUo,3474
+kinfer/export/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+kinfer/export/jax.py,sha256=gZWWUJKb3MPEtfNYeKfPnypeK0elJwGuDYfpCW0GqNY,1413
+kinfer/export/pytorch.py,sha256=rwHGOTcd08IeL5-HZ_7JVUMPMHleiHI8-TkShAoRrYM,1612
+kinfer/export/common.py,sha256=QWcLXY-jzZCbrVPWN7gFGLE5RbnZHLhuuvoH9agEvwk,1244
+kinfer/rust_bindings/Cargo.toml,sha256=i1RGB9VNd9Q4FJ6gGwjZJQYo8DBBvpVWf3GJ95EfVgM,637
+kinfer/rust_bindings/pyproject.toml,sha256=jLcJuHCnQRh9HWR_R7a9qLHwj6LMBgnHyeKK_DruO1Y,135
+kinfer/rust_bindings/rust_bindings.pyi,sha256=9u0_TuHwV9xQH3O5OKbYfUZ_Uj6bsUwEOW0bfZujRMQ,1709
+kinfer/rust_bindings/src/lib.rs,sha256=ydTTyx68d1fZDzc6kjFpU9ePDQa9IBD8Gf2rPy-kutI,10840
+kinfer/rust_bindings/src/bin/stub_gen.rs,sha256=hhoVGnaSfazbSfj5a4x6mPicGPOgWQAfsDmiPej0B6Y,133
+kinfer-0.4.1.dist-info/RECORD,,
+kinfer-0.4.1.dist-info/WHEEL,sha256=rtKwvZSAzWV03G3Ircwq_TRBlj2DQz4ocNXl0bd9DbU,136
+kinfer-0.4.1.dist-info/top_level.txt,sha256=6mY_t3PYr3Dm0dpqMk80uSnArbvGfCFkxOh1QWtgDEo,7
+kinfer-0.4.1.dist-info/METADATA,sha256=yHLwDqimUopYHXEXsUyN_VBAxH2yBKvTxCWX5SbhvPQ,1745
+kinfer-0.4.1.dist-info/licenses/LICENSE,sha256=Qw-Z0XTwS-diSW91e_jLeBPX9zZbAatOJTBLdPHPaC0,1069

kinfer-0.4.0.dist-info/RECORD

@@ -1,26 +0,0 @@
-kinfer/rust_bindings.cpython-311-darwin.so,sha256=lwkKH5R0rx9hsMdjtsNcSWwdFM93Bz2Eg_91Hpcx7vg,1931696
-kinfer/requirements.txt,sha256=j08HO4ptA5afuj99j8FlAP2qla5Zf4_OiEBtgAmF7Jg,90
-kinfer/__init__.py,sha256=YbtJIepEE4pbjYbdgFCV-rBP90AnQpBfDaprkflBmEE,99
-kinfer/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-kinfer/rust_bindings.pyi,sha256=oQKlW_bRw_c7fgVMbqiscPJ57pli25DkTCA_odsMOSU,1639
-kinfer/rust/Cargo.toml,sha256=0SLLbjtoODLrSMcyPdg22Ora-JNdaCqbMKb1c5OQblU,404
-kinfer/rust/src/runtime.rs,sha256=eDflFSIa4IxOmG3SInwrHrYc02UzEANZMkK-liA7pVA,3407
-kinfer/rust/src/lib.rs,sha256=Z3dWdhKhqhWqPJee6vWde4ptqquOYW6W9wB0wyaKCyk,71
-kinfer/rust/src/model.rs,sha256=pODQzD-roXpCJ7-kzaE6uULgqHgsHKb3oQjv-uKtCYk,10668
-kinfer/common/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-kinfer/common/types.py,sha256=Yf3c8Tui7vWguXMmxxa_QynwydCPu9RxrJS8JJMDuKg,147
-kinfer/export/serialize.py,sha256=pOxXdcWMoaNfnAUhgzHUTb7GtZvn0MYtKskIarHzAiY,3007
-kinfer/export/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-kinfer/export/jax.py,sha256=nbZB7Bxed4MQmfBZrH_Ml6PnkHQoyq5jdOT01DxiZDA,1337
-kinfer/export/pytorch.py,sha256=twGpcX6OaJgiiwQ0teo_4rrCFTrCqu-ULH2UEHl8auc,1477
-kinfer/export/common.py,sha256=ZiOAtehb2_ATg7t3tWzxbHUkGhtJ8vqjlaaw94-Iux4,1021
-kinfer/rust_bindings/Cargo.toml,sha256=i1RGB9VNd9Q4FJ6gGwjZJQYo8DBBvpVWf3GJ95EfVgM,637
-kinfer/rust_bindings/pyproject.toml,sha256=jLcJuHCnQRh9HWR_R7a9qLHwj6LMBgnHyeKK_DruO1Y,135
-kinfer/rust_bindings/rust_bindings.pyi,sha256=oQKlW_bRw_c7fgVMbqiscPJ57pli25DkTCA_odsMOSU,1639
-kinfer/rust_bindings/src/lib.rs,sha256=lKYr6Imu0Z-XxDKCguJLfrhIuai6S2bEhRAe__9sFnc,10196
-kinfer/rust_bindings/src/bin/stub_gen.rs,sha256=hhoVGnaSfazbSfj5a4x6mPicGPOgWQAfsDmiPej0B6Y,133
-kinfer-0.4.0.dist-info/RECORD,,
-kinfer-0.4.0.dist-info/WHEEL,sha256=rtKwvZSAzWV03G3Ircwq_TRBlj2DQz4ocNXl0bd9DbU,136
-kinfer-0.4.0.dist-info/top_level.txt,sha256=6mY_t3PYr3Dm0dpqMk80uSnArbvGfCFkxOh1QWtgDEo,7
-kinfer-0.4.0.dist-info/METADATA,sha256=2gGRmUwgmAoYK8kRG5qTtteIMrdaUYW7Edvpk_32yas,1745
-kinfer-0.4.0.dist-info/licenses/LICENSE,sha256=Qw-Z0XTwS-diSW91e_jLeBPX9zZbAatOJTBLdPHPaC0,1069