kinfer 0.4.2__tar.gz → 0.4.3__tar.gz

{kinfer-0.4.2 → kinfer-0.4.3}/Cargo.toml

@@ -8,7 +8,7 @@ resolver = "2"
 
 [workspace.package]
 
-version = "0.4.2"
+version = "0.4.3"
 authors = ["Wesley Maa <wesley@kscale.dev>", "Benjamin Bolte <ben@kscale.dev>"]
 edition = "2021"
 description = "K-Scale Inference Library"

{kinfer-0.4.2 → kinfer-0.4.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: kinfer
-Version: 0.4.2
+Version: 0.4.3
 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
@@ -21,7 +21,7 @@ Provides-Extra: pytorch
 Requires-Dist: torch; extra == "pytorch"
 Provides-Extra: jax
 Requires-Dist: tensorflow; extra == "jax"
-Requires-Dist: tf2onnx; extra == "jax"
+Requires-Dist: tf2onnx>=1.16.0; extra == "jax"
 Requires-Dist: jax; extra == "jax"
 Requires-Dist: equinox; extra == "jax"
 Requires-Dist: numpy<2; extra == "jax"
@@ -34,7 +34,7 @@ Requires-Dist: ruff; extra == "all"
 Requires-Dist: types-tensorflow; extra == "all"
 Requires-Dist: torch; extra == "all"
 Requires-Dist: tensorflow; extra == "all"
-Requires-Dist: tf2onnx; extra == "all"
+Requires-Dist: tf2onnx>=1.16.0; extra == "all"
 Requires-Dist: jax; extra == "all"
 Requires-Dist: equinox; extra == "all"
 Requires-Dist: numpy<2; extra == "all"

{kinfer-0.4.2 → kinfer-0.4.3}/kinfer/export/common.py

@@ -37,5 +37,8 @@ def get_shape(
                 raise ValueError("`carry_shape` must be provided for `carry`")
             return carry_shape
 
+        case "time":
+            return (1,)
+
         case _:
             raise ValueError(f"Unknown tensor name: {name}")

{kinfer-0.4.2 → kinfer-0.4.3}/kinfer/rust/src/model.rs

@@ -47,6 +47,7 @@ pub trait ModelProvider: Send + Sync {
     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_time(&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,
@@ -188,6 +189,15 @@ impl ModelRunner {
                         .into());
                     }
                 }
+                "time" => {
+                    if *dims != vec![1] {
+                        return Err(format!(
+                            "Expected shape [1] for input `{}`, got {:?}",
+                            input.name, dims
+                        )
+                        .into());
+                    }
+                }
                 "carry" => {
                     if dims != carry_shape {
                         return Err(format!(
@@ -269,6 +279,7 @@ impl ModelRunner {
                 "gyroscope" => futures.push(self.provider.get_gyroscope()),
                 "command" => futures.push(self.provider.get_command()),
                 "carry" => futures.push(self.provider.get_carry(carry.clone())),
+                "time" => futures.push(self.provider.get_time()),
                 _ => return Err(format!("Unknown input name: {}", name).into()),
             }
         }

{kinfer-0.4.2 → kinfer-0.4.3}/kinfer/rust/src/runtime.rs

@@ -3,8 +3,8 @@ use std::sync::Arc;
 use tokio::runtime::Runtime;
 
 use crate::model::{ModelError, ModelRunner};
-use std::time::{Duration, Instant};
-use tokio::time::sleep;
+use std::time::Duration;
+use tokio::time::interval;
 
 pub struct ModelRuntime {
     model_runner: Arc<ModelRunner>,
@@ -61,7 +61,10 @@ impl ModelRuntime {
                 .get_joint_angles()
                 .await
                 .map_err(|e| ModelError::Provider(e.to_string()))?;
-            let mut last_time = Instant::now();
+
+            // Wait for the first tick, since it happens immediately.
+            let mut interval = interval(dt);
+            interval.tick().await;
 
             while running.load(Ordering::Relaxed) {
                 let (output, next_carry) = model_runner
@@ -80,10 +83,7 @@ impl ModelRuntime {
                         .take_action(interp_joint_positions * magnitude_factor)
                         .await
                         .map_err(|e| ModelError::Provider(e.to_string()))?;
-                    last_time = last_time + dt;
-                    if let Some(sleep_duration) = last_time.checked_duration_since(Instant::now()) {
-                        sleep(sleep_duration).await;
-                    }
+                    interval.tick().await;
                 }
 
                 joint_positions = output;

{kinfer-0.4.2 → kinfer-0.4.3}/kinfer/rust_bindings/src/lib.rs

@@ -10,7 +10,7 @@ use pyo3_stub_gen::define_stub_info_gatherer;
 use pyo3_stub_gen::derive::{gen_stub_pyclass, gen_stub_pyfunction, gen_stub_pymethods};
 use std::sync::Arc;
 use std::sync::Mutex;
-
+use std::time::Instant;
 #[pyfunction]
 #[gen_stub_pyfunction]
 fn get_version() -> String {
@@ -73,6 +73,10 @@ impl ModelProviderABC {
         Err(PyNotImplementedError::new_err("Must override get_command"))
     }
 
+    fn get_time<'py>(&self) -> PyResult<Bound<'py, PyArray1<f32>>> {
+        Err(PyNotImplementedError::new_err("Must override get_time"))
+    }
+
     fn take_action<'py>(
         &self,
         joint_names: Vec<String>,
@@ -92,13 +96,17 @@ impl ModelProviderABC {
 #[derive(Clone)]
 struct PyModelProvider {
     obj: Arc<Py<ModelProviderABC>>,
+    start_time: Instant,
 }
 
 #[pymethods]
 impl PyModelProvider {
     #[new]
     fn new(obj: Py<ModelProviderABC>) -> Self {
-        Self { obj: Arc::new(obj) }
+        Self {
+            obj: Arc::new(obj),
+            start_time: Instant::now(),
+        }
     }
 }
 
@@ -182,6 +190,18 @@ impl ModelProvider for PyModelProvider {
         Ok(args)
     }
 
+    async fn get_time(&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_time", 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)
     }
@@ -217,9 +237,7 @@ struct PyModelRunner {
 impl PyModelRunner {
     #[new]
     fn new(model_path: String, provider: Py<ModelProviderABC>) -> PyResult<Self> {
-        let input_provider = Arc::new(PyModelProvider {
-            obj: Arc::new(provider),
-        });
+        let input_provider = Arc::new(PyModelProvider::new(provider));
 
         let runner = tokio::runtime::Runtime::new().unwrap().block_on(async {
             ModelRunner::new(model_path, input_provider)

{kinfer-0.4.2 → kinfer-0.4.3}/kinfer.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: kinfer
-Version: 0.4.2
+Version: 0.4.3
 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
@@ -21,7 +21,7 @@ Provides-Extra: pytorch
 Requires-Dist: torch; extra == "pytorch"
 Provides-Extra: jax
 Requires-Dist: tensorflow; extra == "jax"
-Requires-Dist: tf2onnx; extra == "jax"
+Requires-Dist: tf2onnx>=1.16.0; extra == "jax"
 Requires-Dist: jax; extra == "jax"
 Requires-Dist: equinox; extra == "jax"
 Requires-Dist: numpy<2; extra == "jax"
@@ -34,7 +34,7 @@ Requires-Dist: ruff; extra == "all"
 Requires-Dist: types-tensorflow; extra == "all"
 Requires-Dist: torch; extra == "all"
 Requires-Dist: tensorflow; extra == "all"
-Requires-Dist: tf2onnx; extra == "all"
+Requires-Dist: tf2onnx>=1.16.0; extra == "all"
 Requires-Dist: jax; extra == "all"
 Requires-Dist: equinox; extra == "all"
 Requires-Dist: numpy<2; extra == "all"

{kinfer-0.4.2 → kinfer-0.4.3}/kinfer.egg-info/requires.txt

@@ -11,7 +11,7 @@ ruff
 types-tensorflow
 torch
 tensorflow
-tf2onnx
+tf2onnx>=1.16.0
 jax
 equinox
 numpy<2
@@ -26,7 +26,7 @@ types-tensorflow
 
 [jax]
 tensorflow
-tf2onnx
+tf2onnx>=1.16.0
 jax
 equinox
 numpy<2

{kinfer-0.4.2 → kinfer-0.4.3}/setup.py

@@ -39,7 +39,7 @@ requirements_pytorch = [
 
 requirements_jax = [
     "tensorflow",
-    "tf2onnx",
+    "tf2onnx>=1.16.0",
     "jax",
     "equinox",
     "numpy<2",

{kinfer-0.4.2 → kinfer-0.4.3}/tests/test_jax.py

@@ -18,6 +18,7 @@ logger = logging.getLogger(__name__)
 JOINT_NAMES = ["left_arm", "right_arm", "left_leg", "right_leg"]
 NUM_JOINTS = len(JOINT_NAMES)
 CARRY_SIZE = 10
+NUM_COMMANDS = 4
 
 
 @jax.jit
@@ -32,6 +33,8 @@ def step_fn(
     projected_gravity: jnp.ndarray,
     accelerometer: jnp.ndarray,
     gyroscope: jnp.ndarray,
+    command: jnp.ndarray,
+    time: jnp.ndarray,
     carry: jnp.ndarray,
 ) -> tuple[jnp.ndarray, jnp.ndarray]:
     output = (
@@ -40,6 +43,9 @@ def step_fn(
         + projected_gravity.mean()
         + accelerometer.mean()
         + gyroscope.mean()
+        + command.mean()
+        + jnp.cos(time).mean()
+        + jnp.sin(time).mean()
         + carry.mean()
     ) * joint_angles
     next_carry = carry + 1
@@ -54,6 +60,7 @@ def test_export(tmpdir: Path) -> None:
     step_fn_onnx = export_fn(
         model=step_fn,
         num_joints=NUM_JOINTS,
+        num_commands=NUM_COMMANDS,
         carry_shape=(CARRY_SIZE,),
     )
 
@@ -61,6 +68,7 @@ def test_export(tmpdir: Path) -> None:
         init_fn_onnx,
         step_fn_onnx,
         joint_names=JOINT_NAMES,
+        num_commands=NUM_COMMANDS,
         carry_shape=(CARRY_SIZE,),
     )
 
@@ -88,6 +96,12 @@ def test_export(tmpdir: Path) -> None:
         def get_gyroscope(self) -> np.ndarray:
             return np.random.randn(3)
 
+        def get_command(self) -> np.ndarray:
+            return np.random.randn(NUM_COMMANDS)
+
+        def get_time(self) -> np.ndarray:
+            return np.random.randn(1)
+
         def take_action(self, joint_names: Sequence[str], action: np.ndarray) -> None:
             assert joint_names == JOINT_NAMES
             assert action.shape == (NUM_JOINTS,)

{kinfer-0.4.2 → kinfer-0.4.3}/tests/test_pytorch.py

@@ -38,6 +38,7 @@ def step_fn(
     accelerometer: Tensor,
     gyroscope: Tensor,
     command: Tensor,
+    time: Tensor,
     carry: Tensor,
 ) -> tuple[Tensor, Tensor]:
     output = (
@@ -47,6 +48,8 @@ def step_fn(
         + accelerometer.mean()
         + gyroscope.mean()
         + command.mean()
+        + torch.cos(time).mean()
+        + torch.sin(time).mean()
         + carry.mean()
     ) * joint_angles
     next_carry = carry + 1
@@ -120,6 +123,9 @@ def test_export(tmpdir: Path) -> None:
         def get_command(self) -> np.ndarray:
             return np.random.randn(NUM_COMMANDS)
 
+        def get_time(self) -> np.ndarray:
+            return np.random.randn(1)
+
         def take_action(self, joint_names: Sequence[str], action: np.ndarray) -> None:
             assert joint_names == JOINT_NAMES
             assert action.shape == (NUM_JOINTS,)