imt-ring 1.2.2__tar.gz → 1.3.1__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- {imt_ring-1.2.2 → imt_ring-1.3.1}/PKG-INFO +1 -1
- {imt_ring-1.2.2 → imt_ring-1.3.1}/pyproject.toml +1 -1
- {imt_ring-1.2.2 → imt_ring-1.3.1}/src/imt_ring.egg-info/PKG-INFO +1 -1
- {imt_ring-1.2.2 → imt_ring-1.3.1}/src/ring/algorithms/generator/base.py +9 -15
- {imt_ring-1.2.2 → imt_ring-1.3.1}/src/ring/algorithms/generator/batch.py +33 -11
- {imt_ring-1.2.2 → imt_ring-1.3.1}/src/ring/algorithms/jcalc.py +112 -0
- {imt_ring-1.2.2 → imt_ring-1.3.1}/src/ring/base.py +6 -7
- {imt_ring-1.2.2 → imt_ring-1.3.1}/src/ring/io/examples/exclude/standard_sys.xml +40 -40
- {imt_ring-1.2.2 → imt_ring-1.3.1}/src/ring/io/examples/exclude/standard_sys_rr_imp.xml +40 -40
- imt_ring-1.3.1/src/ring/ml/__init__.py +63 -0
- {imt_ring-1.2.2 → imt_ring-1.3.1}/src/ring/ml/base.py +3 -3
- {imt_ring-1.2.2 → imt_ring-1.3.1}/src/ring/ml/callbacks.py +4 -0
- {imt_ring-1.2.2 → imt_ring-1.3.1}/src/ring/rendering/base_render.py +12 -7
- {imt_ring-1.2.2 → imt_ring-1.3.1}/src/ring/utils/__init__.py +1 -0
- {imt_ring-1.2.2 → imt_ring-1.3.1}/src/ring/utils/utils.py +14 -0
- {imt_ring-1.2.2 → imt_ring-1.3.1}/tests/test_jcalc.py +5 -0
- {imt_ring-1.2.2 → imt_ring-1.3.1}/tests/test_randomize.py +28 -28
- {imt_ring-1.2.2 → imt_ring-1.3.1}/tests/test_sys_composer.py +7 -7
- {imt_ring-1.2.2 → imt_ring-1.3.1}/tests/test_train.py +30 -3
- imt_ring-1.2.2/src/ring/ml/__init__.py +0 -33
- {imt_ring-1.2.2 → imt_ring-1.3.1}/readme.md +0 -0
- {imt_ring-1.2.2 → imt_ring-1.3.1}/setup.cfg +0 -0
- {imt_ring-1.2.2 → imt_ring-1.3.1}/src/imt_ring.egg-info/SOURCES.txt +0 -0
- {imt_ring-1.2.2 → imt_ring-1.3.1}/src/imt_ring.egg-info/dependency_links.txt +0 -0
- {imt_ring-1.2.2 → imt_ring-1.3.1}/src/imt_ring.egg-info/requires.txt +0 -0
- {imt_ring-1.2.2 → imt_ring-1.3.1}/src/imt_ring.egg-info/top_level.txt +0 -0
- {imt_ring-1.2.2 → imt_ring-1.3.1}/src/ring/__init__.py +0 -0
- {imt_ring-1.2.2 → imt_ring-1.3.1}/src/ring/algebra.py +0 -0
- {imt_ring-1.2.2 → imt_ring-1.3.1}/src/ring/algorithms/__init__.py +0 -0
- {imt_ring-1.2.2 → imt_ring-1.3.1}/src/ring/algorithms/_random.py +0 -0
- {imt_ring-1.2.2 → imt_ring-1.3.1}/src/ring/algorithms/custom_joints/__init__.py +0 -0
- {imt_ring-1.2.2 → imt_ring-1.3.1}/src/ring/algorithms/custom_joints/rr_imp_joint.py +0 -0
- {imt_ring-1.2.2 → imt_ring-1.3.1}/src/ring/algorithms/custom_joints/rr_joint.py +0 -0
- {imt_ring-1.2.2 → imt_ring-1.3.1}/src/ring/algorithms/custom_joints/suntay.py +0 -0
- {imt_ring-1.2.2 → imt_ring-1.3.1}/src/ring/algorithms/dynamics.py +0 -0
- {imt_ring-1.2.2 → imt_ring-1.3.1}/src/ring/algorithms/generator/__init__.py +0 -0
- {imt_ring-1.2.2 → imt_ring-1.3.1}/src/ring/algorithms/generator/motion_artifacts.py +0 -0
- {imt_ring-1.2.2 → imt_ring-1.3.1}/src/ring/algorithms/generator/pd_control.py +0 -0
- {imt_ring-1.2.2 → imt_ring-1.3.1}/src/ring/algorithms/generator/randomize.py +0 -0
- {imt_ring-1.2.2 → imt_ring-1.3.1}/src/ring/algorithms/generator/transforms.py +0 -0
- {imt_ring-1.2.2 → imt_ring-1.3.1}/src/ring/algorithms/generator/types.py +0 -0
- {imt_ring-1.2.2 → imt_ring-1.3.1}/src/ring/algorithms/kinematics.py +0 -0
- {imt_ring-1.2.2 → imt_ring-1.3.1}/src/ring/algorithms/sensors.py +0 -0
- {imt_ring-1.2.2 → imt_ring-1.3.1}/src/ring/io/__init__.py +0 -0
- {imt_ring-1.2.2 → imt_ring-1.3.1}/src/ring/io/examples/branched.xml +0 -0
- {imt_ring-1.2.2 → imt_ring-1.3.1}/src/ring/io/examples/exclude/knee_trans_dof.xml +0 -0
- {imt_ring-1.2.2 → imt_ring-1.3.1}/src/ring/io/examples/inv_pendulum.xml +0 -0
- {imt_ring-1.2.2 → imt_ring-1.3.1}/src/ring/io/examples/knee_flexible_imus.xml +0 -0
- {imt_ring-1.2.2 → imt_ring-1.3.1}/src/ring/io/examples/spherical_stiff.xml +0 -0
- {imt_ring-1.2.2 → imt_ring-1.3.1}/src/ring/io/examples/symmetric.xml +0 -0
- {imt_ring-1.2.2 → imt_ring-1.3.1}/src/ring/io/examples/test_all_1.xml +0 -0
- {imt_ring-1.2.2 → imt_ring-1.3.1}/src/ring/io/examples/test_all_2.xml +0 -0
- {imt_ring-1.2.2 → imt_ring-1.3.1}/src/ring/io/examples/test_ang0_pos0.xml +0 -0
- {imt_ring-1.2.2 → imt_ring-1.3.1}/src/ring/io/examples/test_control.xml +0 -0
- {imt_ring-1.2.2 → imt_ring-1.3.1}/src/ring/io/examples/test_double_pendulum.xml +0 -0
- {imt_ring-1.2.2 → imt_ring-1.3.1}/src/ring/io/examples/test_free.xml +0 -0
- {imt_ring-1.2.2 → imt_ring-1.3.1}/src/ring/io/examples/test_kinematics.xml +0 -0
- {imt_ring-1.2.2 → imt_ring-1.3.1}/src/ring/io/examples/test_morph_system/four_seg_seg1.xml +0 -0
- {imt_ring-1.2.2 → imt_ring-1.3.1}/src/ring/io/examples/test_morph_system/four_seg_seg3.xml +0 -0
- {imt_ring-1.2.2 → imt_ring-1.3.1}/src/ring/io/examples/test_randomize_position.xml +0 -0
- {imt_ring-1.2.2 → imt_ring-1.3.1}/src/ring/io/examples/test_sensors.xml +0 -0
- {imt_ring-1.2.2 → imt_ring-1.3.1}/src/ring/io/examples/test_three_seg_seg2.xml +0 -0
- {imt_ring-1.2.2 → imt_ring-1.3.1}/src/ring/io/examples.py +0 -0
- {imt_ring-1.2.2 → imt_ring-1.3.1}/src/ring/io/test_examples.py +0 -0
- {imt_ring-1.2.2 → imt_ring-1.3.1}/src/ring/io/xml/__init__.py +0 -0
- {imt_ring-1.2.2 → imt_ring-1.3.1}/src/ring/io/xml/abstract.py +0 -0
- {imt_ring-1.2.2 → imt_ring-1.3.1}/src/ring/io/xml/from_xml.py +0 -0
- {imt_ring-1.2.2 → imt_ring-1.3.1}/src/ring/io/xml/test_from_xml.py +0 -0
- {imt_ring-1.2.2 → imt_ring-1.3.1}/src/ring/io/xml/test_to_xml.py +0 -0
- {imt_ring-1.2.2 → imt_ring-1.3.1}/src/ring/io/xml/to_xml.py +0 -0
- {imt_ring-1.2.2 → imt_ring-1.3.1}/src/ring/maths.py +0 -0
- {imt_ring-1.2.2 → imt_ring-1.3.1}/src/ring/ml/ml_utils.py +0 -0
- {imt_ring-1.2.2 → imt_ring-1.3.1}/src/ring/ml/optimizer.py +0 -0
- {imt_ring-1.2.2 → imt_ring-1.3.1}/src/ring/ml/params/0x13e3518065c21cd8.pickle +0 -0
- {imt_ring-1.2.2 → imt_ring-1.3.1}/src/ring/ml/ringnet.py +0 -0
- {imt_ring-1.2.2 → imt_ring-1.3.1}/src/ring/ml/train.py +0 -0
- {imt_ring-1.2.2 → imt_ring-1.3.1}/src/ring/ml/training_loop.py +0 -0
- {imt_ring-1.2.2 → imt_ring-1.3.1}/src/ring/rendering/__init__.py +0 -0
- {imt_ring-1.2.2 → imt_ring-1.3.1}/src/ring/rendering/mujoco_render.py +0 -0
- {imt_ring-1.2.2 → imt_ring-1.3.1}/src/ring/rendering/vispy_render.py +0 -0
- {imt_ring-1.2.2 → imt_ring-1.3.1}/src/ring/rendering/vispy_visuals.py +0 -0
- {imt_ring-1.2.2 → imt_ring-1.3.1}/src/ring/sim2real/__init__.py +0 -0
- {imt_ring-1.2.2 → imt_ring-1.3.1}/src/ring/sim2real/sim2real.py +0 -0
- {imt_ring-1.2.2 → imt_ring-1.3.1}/src/ring/spatial.py +0 -0
- {imt_ring-1.2.2 → imt_ring-1.3.1}/src/ring/sys_composer/__init__.py +0 -0
- {imt_ring-1.2.2 → imt_ring-1.3.1}/src/ring/sys_composer/delete_sys.py +0 -0
- {imt_ring-1.2.2 → imt_ring-1.3.1}/src/ring/sys_composer/inject_sys.py +0 -0
- {imt_ring-1.2.2 → imt_ring-1.3.1}/src/ring/sys_composer/morph_sys.py +0 -0
- {imt_ring-1.2.2 → imt_ring-1.3.1}/src/ring/utils/batchsize.py +0 -0
- {imt_ring-1.2.2 → imt_ring-1.3.1}/src/ring/utils/colab.py +0 -0
- {imt_ring-1.2.2 → imt_ring-1.3.1}/src/ring/utils/hdf5.py +0 -0
- {imt_ring-1.2.2 → imt_ring-1.3.1}/src/ring/utils/normalizer.py +0 -0
- {imt_ring-1.2.2 → imt_ring-1.3.1}/src/ring/utils/path.py +0 -0
- {imt_ring-1.2.2 → imt_ring-1.3.1}/tests/test_algebra.py +0 -0
- {imt_ring-1.2.2 → imt_ring-1.3.1}/tests/test_base.py +0 -0
- {imt_ring-1.2.2 → imt_ring-1.3.1}/tests/test_custom_joints.py +0 -0
- {imt_ring-1.2.2 → imt_ring-1.3.1}/tests/test_dynamics.py +0 -0
- {imt_ring-1.2.2 → imt_ring-1.3.1}/tests/test_generator.py +0 -0
- {imt_ring-1.2.2 → imt_ring-1.3.1}/tests/test_jit.py +0 -0
- {imt_ring-1.2.2 → imt_ring-1.3.1}/tests/test_kinematics.py +0 -0
- {imt_ring-1.2.2 → imt_ring-1.3.1}/tests/test_maths.py +0 -0
- {imt_ring-1.2.2 → imt_ring-1.3.1}/tests/test_ml_utils.py +0 -0
- {imt_ring-1.2.2 → imt_ring-1.3.1}/tests/test_motion_artifacts.py +0 -0
- {imt_ring-1.2.2 → imt_ring-1.3.1}/tests/test_pd_control.py +0 -0
- {imt_ring-1.2.2 → imt_ring-1.3.1}/tests/test_random.py +0 -0
- {imt_ring-1.2.2 → imt_ring-1.3.1}/tests/test_rcmg.py +0 -0
- {imt_ring-1.2.2 → imt_ring-1.3.1}/tests/test_render.py +0 -0
- {imt_ring-1.2.2 → imt_ring-1.3.1}/tests/test_sensors.py +0 -0
- {imt_ring-1.2.2 → imt_ring-1.3.1}/tests/test_sim2real.py +0 -0
- {imt_ring-1.2.2 → imt_ring-1.3.1}/tests/test_utils.py +0 -0
@@ -4,6 +4,8 @@ import warnings
|
|
4
4
|
|
5
5
|
import jax
|
6
6
|
import jax.numpy as jnp
|
7
|
+
import tree_utils
|
8
|
+
|
7
9
|
from ring import base
|
8
10
|
from ring import utils
|
9
11
|
from ring.algorithms import jcalc
|
@@ -13,7 +15,6 @@ from ring.algorithms.generator import motion_artifacts
|
|
13
15
|
from ring.algorithms.generator import randomize
|
14
16
|
from ring.algorithms.generator import transforms
|
15
17
|
from ring.algorithms.generator import types
|
16
|
-
import tree_utils
|
17
18
|
|
18
19
|
|
19
20
|
class RCMG:
|
@@ -108,23 +109,20 @@ class RCMG:
|
|
108
109
|
partial_build_gen(sys=_sys, config=_config, sys_ml=sys_ml)
|
109
110
|
)
|
110
111
|
|
111
|
-
def _to_data(self, sizes, seed
|
112
|
-
return batch.batch_generators_eager_to_list(
|
113
|
-
self.gens, sizes, seed=seed, jit=jit
|
114
|
-
)
|
112
|
+
def _to_data(self, sizes, seed):
|
113
|
+
return batch.batch_generators_eager_to_list(self.gens, sizes, seed=seed)
|
115
114
|
|
116
|
-
def to_list(self, sizes: int | list[int] = 1, seed: int = 1
|
117
|
-
return self._to_data(sizes, seed
|
115
|
+
def to_list(self, sizes: int | list[int] = 1, seed: int = 1):
|
116
|
+
return self._to_data(sizes, seed)
|
118
117
|
|
119
118
|
def to_pickle(
|
120
119
|
self,
|
121
120
|
path: str,
|
122
121
|
sizes: int | list[int] = 1,
|
123
122
|
seed: int = 1,
|
124
|
-
jit: bool = False,
|
125
123
|
overwrite: bool = True,
|
126
124
|
) -> None:
|
127
|
-
data = tree_utils.tree_batch(self._to_data(sizes, seed
|
125
|
+
data = tree_utils.tree_batch(self._to_data(sizes, seed))
|
128
126
|
utils.pickle_save(data, path, overwrite=overwrite)
|
129
127
|
|
130
128
|
def to_hdf5(
|
@@ -132,10 +130,9 @@ class RCMG:
|
|
132
130
|
path: str,
|
133
131
|
sizes: int | list[int] = 1,
|
134
132
|
seed: int = 1,
|
135
|
-
jit: bool = False,
|
136
133
|
overwrite: bool = True,
|
137
134
|
) -> None:
|
138
|
-
data = tree_utils.tree_batch(self._to_data(sizes, seed
|
135
|
+
data = tree_utils.tree_batch(self._to_data(sizes, seed))
|
139
136
|
utils.hdf5_save(path, data, overwrite=overwrite)
|
140
137
|
|
141
138
|
def to_eager_gen(
|
@@ -143,11 +140,8 @@ class RCMG:
|
|
143
140
|
batchsize: int = 1,
|
144
141
|
sizes: int | list[int] = 1,
|
145
142
|
seed: int = 1,
|
146
|
-
jit: bool = False,
|
147
143
|
) -> types.BatchedGenerator:
|
148
|
-
return batch.batch_generators_eager(
|
149
|
-
self.gens, sizes, batchsize, seed=seed, jit=jit
|
150
|
-
)
|
144
|
+
return batch.batch_generators_eager(self.gens, sizes, batchsize, seed=seed)
|
151
145
|
|
152
146
|
def to_lazy_gen(
|
153
147
|
self, sizes: int | list[int] = 1, jit: bool = True
|
@@ -6,12 +6,13 @@ import warnings
|
|
6
6
|
import jax
|
7
7
|
import jax.numpy as jnp
|
8
8
|
import numpy as np
|
9
|
-
from ring import utils
|
10
|
-
from ring.algorithms.generator import types
|
11
9
|
from tqdm import tqdm
|
12
10
|
import tree_utils
|
13
11
|
from tree_utils import tree_batch
|
14
12
|
|
13
|
+
from ring import utils
|
14
|
+
from ring.algorithms.generator import types
|
15
|
+
|
15
16
|
|
16
17
|
def _build_batch_matrix(batchsizes: list[int]) -> jax.Array:
|
17
18
|
arr = []
|
@@ -61,11 +62,24 @@ def batch_generators_lazy(
|
|
61
62
|
return generator
|
62
63
|
|
63
64
|
|
65
|
+
def _number_of_executions_required(size: int) -> int:
|
66
|
+
vmap_threshold = 128
|
67
|
+
_, vmap = utils.distribute_batchsize(size)
|
68
|
+
|
69
|
+
primes = iter(utils.primes(vmap))
|
70
|
+
n_calls = 1
|
71
|
+
while vmap > vmap_threshold:
|
72
|
+
prime = next(primes)
|
73
|
+
n_calls *= prime
|
74
|
+
vmap /= prime
|
75
|
+
|
76
|
+
return n_calls
|
77
|
+
|
78
|
+
|
64
79
|
def batch_generators_eager_to_list(
|
65
80
|
generators: types.Generator | list[types.Generator],
|
66
81
|
sizes: int | list[int],
|
67
82
|
seed: int = 1,
|
68
|
-
jit: bool = True,
|
69
83
|
) -> list[tree_utils.PyTree]:
|
70
84
|
"Returns list of unbatched sequences as numpy arrays."
|
71
85
|
generators, sizes = _process_sizes_batchsizes_generators(generators, sizes)
|
@@ -73,11 +87,20 @@ def batch_generators_eager_to_list(
|
|
73
87
|
key = jax.random.PRNGKey(seed)
|
74
88
|
data = []
|
75
89
|
for gen, size in tqdm(zip(generators, sizes), desc="eager data generation"):
|
76
|
-
|
77
|
-
|
78
|
-
#
|
79
|
-
|
80
|
-
|
90
|
+
|
91
|
+
n_calls = _number_of_executions_required(size)
|
92
|
+
# decrease size by n_calls times
|
93
|
+
size = int(size / n_calls)
|
94
|
+
jit = True if n_calls > 1 else False
|
95
|
+
gen_jit = batch_generators_lazy(gen, size, jit=jit)
|
96
|
+
|
97
|
+
for _ in range(n_calls):
|
98
|
+
key, consume = jax.random.split(key)
|
99
|
+
sample = gen_jit(consume)
|
100
|
+
# converts also to numpy
|
101
|
+
sample = jax.device_get(sample)
|
102
|
+
data.extend([jax.tree_map(lambda a: a[i], sample) for i in range(size)])
|
103
|
+
|
81
104
|
return data
|
82
105
|
|
83
106
|
|
@@ -243,12 +266,11 @@ def batch_generators_eager(
|
|
243
266
|
shuffle: bool = True,
|
244
267
|
drop_last: bool = True,
|
245
268
|
seed: int = 1,
|
246
|
-
jit: bool = True,
|
247
269
|
) -> types.BatchedGenerator:
|
248
270
|
"""Eagerly create a large precomputed generator by calling multiple generators
|
249
271
|
and stacking their output."""
|
250
272
|
|
251
|
-
data = batch_generators_eager_to_list(generators, sizes, seed=seed
|
273
|
+
data = batch_generators_eager_to_list(generators, sizes, seed=seed)
|
252
274
|
return batched_generator_from_list(data, batchsize, shuffle, drop_last)
|
253
275
|
|
254
276
|
|
@@ -270,7 +292,7 @@ def _process_sizes_batchsizes_generators(
|
|
270
292
|
|
271
293
|
assert len(generators) == len(list_sizes)
|
272
294
|
|
273
|
-
_WARN_SIZE =
|
295
|
+
_WARN_SIZE = 1e6 # disable this warning
|
274
296
|
for size in list_sizes:
|
275
297
|
if size >= _WARN_SIZE:
|
276
298
|
warnings.warn(
|
@@ -88,6 +88,118 @@ class MotionConfig:
|
|
88
88
|
assert nomotion_config.is_feasible()
|
89
89
|
return nomotion_config
|
90
90
|
|
91
|
+
@staticmethod
|
92
|
+
def from_register(name: str) -> "MotionConfig":
|
93
|
+
return _registered_motion_configs[name]
|
94
|
+
|
95
|
+
|
96
|
+
_registered_motion_configs = {
|
97
|
+
"hinUndHer": MotionConfig(
|
98
|
+
t_min=0.3,
|
99
|
+
t_max=1.5,
|
100
|
+
dang_max=3.0,
|
101
|
+
delta_ang_min=0.5,
|
102
|
+
pos_min=-1.5,
|
103
|
+
pos_max=1.5,
|
104
|
+
randomized_interpolation_angle=True,
|
105
|
+
cor=True,
|
106
|
+
),
|
107
|
+
"langsam": MotionConfig(
|
108
|
+
t_min=0.2,
|
109
|
+
t_max=1.25,
|
110
|
+
dang_max=2.0,
|
111
|
+
randomized_interpolation_angle=True,
|
112
|
+
dang_max_free_spherical=2.0,
|
113
|
+
cdf_bins_min=1,
|
114
|
+
cdf_bins_max=3,
|
115
|
+
pos_min=-1.5,
|
116
|
+
pos_max=1.5,
|
117
|
+
cor=True,
|
118
|
+
),
|
119
|
+
"standard": MotionConfig(
|
120
|
+
randomized_interpolation_angle=True,
|
121
|
+
cdf_bins_min=1,
|
122
|
+
cdf_bins_max=5,
|
123
|
+
cor=True,
|
124
|
+
),
|
125
|
+
"expFast": MotionConfig(
|
126
|
+
t_min=0.4,
|
127
|
+
t_max=1.1,
|
128
|
+
dang_max=jnp.deg2rad(180),
|
129
|
+
delta_ang_min=jnp.deg2rad(60),
|
130
|
+
delta_ang_max=jnp.deg2rad(110),
|
131
|
+
pos_min=-1.5,
|
132
|
+
pos_max=1.5,
|
133
|
+
range_of_motion_hinge_method="sigmoid",
|
134
|
+
randomized_interpolation_angle=True,
|
135
|
+
cdf_bins_min=1,
|
136
|
+
cdf_bins_max=3,
|
137
|
+
cor=True,
|
138
|
+
),
|
139
|
+
"expSlow": MotionConfig(
|
140
|
+
t_min=0.75,
|
141
|
+
t_max=3.0,
|
142
|
+
dang_min=0.1,
|
143
|
+
dang_max=1.0,
|
144
|
+
dang_min_free_spherical=0.1,
|
145
|
+
delta_ang_min=0.4,
|
146
|
+
dang_max_free_spherical=1.0,
|
147
|
+
delta_ang_max_free_spherical=1.0,
|
148
|
+
dpos_max=0.3,
|
149
|
+
cor_dpos_max=0.3,
|
150
|
+
range_of_motion_hinge_method="sigmoid",
|
151
|
+
randomized_interpolation_angle=True,
|
152
|
+
cdf_bins_min=1,
|
153
|
+
cdf_bins_max=5,
|
154
|
+
cor=True,
|
155
|
+
),
|
156
|
+
"expFastNoSig": MotionConfig(
|
157
|
+
t_min=0.4,
|
158
|
+
t_max=1.1,
|
159
|
+
dang_max=jnp.deg2rad(180),
|
160
|
+
delta_ang_min=jnp.deg2rad(60),
|
161
|
+
delta_ang_max=jnp.deg2rad(110),
|
162
|
+
pos_min=-1.5,
|
163
|
+
pos_max=1.5,
|
164
|
+
randomized_interpolation_angle=True,
|
165
|
+
cdf_bins_min=1,
|
166
|
+
cdf_bins_max=3,
|
167
|
+
cor=True,
|
168
|
+
),
|
169
|
+
"expSlowNoSig": MotionConfig(
|
170
|
+
t_min=0.75,
|
171
|
+
t_max=3.0,
|
172
|
+
dang_min=0.1,
|
173
|
+
dang_max=1.0,
|
174
|
+
dang_min_free_spherical=0.1,
|
175
|
+
delta_ang_min=0.4,
|
176
|
+
dang_max_free_spherical=1.0,
|
177
|
+
delta_ang_max_free_spherical=1.0,
|
178
|
+
dpos_max=0.3,
|
179
|
+
cor_dpos_max=0.3,
|
180
|
+
randomized_interpolation_angle=True,
|
181
|
+
cdf_bins_min=1,
|
182
|
+
cdf_bins_max=3,
|
183
|
+
cor=True,
|
184
|
+
),
|
185
|
+
"verySlow": MotionConfig(
|
186
|
+
t_min=1.5,
|
187
|
+
t_max=5.0,
|
188
|
+
dang_min=jnp.deg2rad(1),
|
189
|
+
dang_max=jnp.deg2rad(30),
|
190
|
+
delta_ang_min=jnp.deg2rad(20),
|
191
|
+
dang_min_free_spherical=jnp.deg2rad(1),
|
192
|
+
dang_max_free_spherical=jnp.deg2rad(10),
|
193
|
+
delta_ang_min_free_spherical=jnp.deg2rad(5),
|
194
|
+
dpos_max=0.3,
|
195
|
+
cor_dpos_max=0.3,
|
196
|
+
randomized_interpolation_angle=True,
|
197
|
+
cdf_bins_min=1,
|
198
|
+
cdf_bins_max=3,
|
199
|
+
cor=True,
|
200
|
+
),
|
201
|
+
}
|
202
|
+
|
91
203
|
|
92
204
|
def _is_feasible_config1(c: MotionConfig) -> bool:
|
93
205
|
t_min, t_max = c.t_min, _to_float(c.t_max, 0.0)
|
@@ -99,15 +99,15 @@ class _Base:
|
|
99
99
|
def ndim(self):
|
100
100
|
return tu.tree_ndim(self)
|
101
101
|
|
102
|
-
def shape(self, axis=0) -> int:
|
103
|
-
|
104
|
-
|
105
|
-
def __len__(self) -> int:
|
106
|
-
Bs = tree_map(lambda arr: arr.shape[0], self)
|
102
|
+
def shape(self, axis: int = 0) -> int:
|
103
|
+
Bs = tree_map(lambda arr: arr.shape[axis], self)
|
107
104
|
Bs = set(jax.tree_util.tree_flatten(Bs)[0])
|
108
105
|
assert len(Bs) == 1
|
109
106
|
return list(Bs)[0]
|
110
107
|
|
108
|
+
def __len__(self) -> int:
|
109
|
+
return self.shape(axis=0)
|
110
|
+
|
111
111
|
|
112
112
|
@struct.dataclass
|
113
113
|
class Transform(_Base):
|
@@ -685,14 +685,13 @@ class System(_Base):
|
|
685
685
|
self,
|
686
686
|
xs: Transform | list[Transform],
|
687
687
|
yhat: dict | jax.Array | np.ndarray,
|
688
|
-
stepframe: int = 1,
|
689
688
|
# by default we don't predict the global rotation
|
690
689
|
transparent_segment_to_root: bool = True,
|
691
690
|
**kwargs,
|
692
691
|
):
|
693
692
|
"`xs` matches `sys`. `yhat` matches `sys_noimu`. `yhat` are child-to-parent."
|
694
693
|
return ring.rendering.render_prediction(
|
695
|
-
self, xs, yhat,
|
694
|
+
self, xs, yhat, transparent_segment_to_root, **kwargs
|
696
695
|
)
|
697
696
|
|
698
697
|
def delete_system(self, link_name: str | list[str], strict: bool = True):
|
@@ -2,101 +2,101 @@
|
|
2
2
|
<x_xy model="arm_1Seg">
|
3
3
|
<options dt="0.01" gravity="0.0 0.0 9.81"/>
|
4
4
|
<worldbody>
|
5
|
-
<body joint="free" name="
|
5
|
+
<body joint="free" name="seg3_1Seg" pos="0.4 0.0 0.0" pos_min="0.2 -0.05 -0.05" pos_max="0.55 0.05 0.05" damping="5.0 5.0 5.0 25.0 25.0 25.0">
|
6
6
|
<geom pos="0.1 0.0 0.0" mass="1.0" color="dustin_exp_blue" edge_color="black" type="box" dim="0.2 0.05 0.05"/>
|
7
7
|
<geom pos="0.05 0.05 0.0" mass="0.1" color="black" edge_color="black" type="box" dim="0.01 0.1 0.01"/>
|
8
8
|
<geom pos="0.15 -0.05 0.0" mass="0.1" color="black" edge_color="black" type="box" dim="0.01 0.1 0.01"/>
|
9
|
-
<omc pos="0.0 0.0 -0.02" name="
|
10
|
-
<body joint="frozen" name="
|
9
|
+
<omc pos="0.0 0.0 -0.02" name="seg3" pos_marker="1"/>
|
10
|
+
<body joint="frozen" name="imu3_1Seg" pos="0.099999994 0.0 0.035" pos_min="0.050000012 -0.05 -0.05" pos_max="0.15 0.05 0.05">
|
11
11
|
<geom mass="0.1" color="dustin_exp_orange" edge_color="black" type="box" dim="0.05 0.03 0.02"/>
|
12
|
-
<omc pos="0.1 0.0 0.015" name="
|
12
|
+
<omc pos="0.1 0.0 0.015" name="seg3" pos_marker="1"/>
|
13
13
|
</body>
|
14
14
|
</body>
|
15
|
-
<body joint="free" name="
|
15
|
+
<body joint="free" name="seg3_2Seg" pos="0.4 0.0 0.0" pos_min="0.2 -0.05 -0.05" pos_max="0.55 0.05 0.05" damping="5.0 5.0 5.0 25.0 25.0 25.0">
|
16
16
|
<geom pos="0.1 0.0 0.0" mass="1.0" color="dustin_exp_blue" edge_color="black" type="box" dim="0.2 0.05 0.05"/>
|
17
17
|
<geom pos="0.05 0.05 0.0" mass="0.1" color="black" edge_color="black" type="box" dim="0.01 0.1 0.01"/>
|
18
18
|
<geom pos="0.15 -0.05 0.0" mass="0.1" color="black" edge_color="black" type="box" dim="0.01 0.1 0.01"/>
|
19
|
-
<omc pos="0.0 0.0 -0.02" name="
|
20
|
-
<body joint="frozen" name="
|
19
|
+
<omc pos="0.0 0.0 -0.02" name="seg3" pos_marker="1"/>
|
20
|
+
<body joint="frozen" name="imu3_2Seg" pos="0.099999994 0.0 0.035" pos_min="0.050000012 -0.05 -0.05" pos_max="0.15 0.05 0.05">
|
21
21
|
<geom mass="0.1" color="dustin_exp_orange" edge_color="black" type="box" dim="0.05 0.03 0.02"/>
|
22
|
-
<omc pos="0.1 0.0 0.015" name="
|
22
|
+
<omc pos="0.1 0.0 0.015" name="seg3" pos_marker="1"/>
|
23
23
|
</body>
|
24
|
-
<body joint="ry" name="
|
24
|
+
<body joint="ry" name="seg4_2Seg" pos="0.20000002 0.0 0.0" pos_min="0.0 -0.05 -0.05" pos_max="0.35 0.05 0.05" damping="3.0">
|
25
25
|
<geom pos="0.1 0.0 0.0" mass="1.0" color="dustin_exp_white" edge_color="black" type="box" dim="0.2 0.05 0.05"/>
|
26
26
|
<geom pos="0.1 0.05 0.0" mass="0.1" color="black" edge_color="black" type="box" dim="0.01 0.1 0.01"/>
|
27
27
|
<geom pos="0.15 -0.05 0.0" mass="0.1" color="black" edge_color="black" type="box" dim="0.01 0.1 0.01"/>
|
28
|
-
<omc pos="0.0 0.0 -0.02" name="
|
29
|
-
<body joint="frozen" name="
|
28
|
+
<omc pos="0.0 0.0 -0.02" name="seg4" pos_marker="2"/>
|
29
|
+
<body joint="frozen" name="imu4_2Seg" pos="0.100000024 0.0 0.035" pos_min="0.050000012 -0.05 -0.05" pos_max="0.14999998 0.05 0.05">
|
30
30
|
<geom mass="0.1" color="dustin_exp_orange" edge_color="black" type="box" dim="0.05 0.03 0.02"/>
|
31
|
-
<omc pos="0.1 0.0 0.015" name="
|
31
|
+
<omc pos="0.1 0.0 0.015" name="seg4" pos_marker="2"/>
|
32
32
|
</body>
|
33
33
|
</body>
|
34
34
|
</body>
|
35
|
-
<body joint="free" name="
|
35
|
+
<body joint="free" name="seg3_3Seg" pos="0.4 0.0 0.0" pos_min="0.2 -0.05 -0.05" pos_max="0.55 0.05 0.05" damping="5.0 5.0 5.0 25.0 25.0 25.0">
|
36
36
|
<geom pos="0.1 0.0 0.0" mass="1.0" color="dustin_exp_blue" edge_color="black" type="box" dim="0.2 0.05 0.05"/>
|
37
37
|
<geom pos="0.05 0.05 0.0" mass="0.1" color="black" edge_color="black" type="box" dim="0.01 0.1 0.01"/>
|
38
38
|
<geom pos="0.15 -0.05 0.0" mass="0.1" color="black" edge_color="black" type="box" dim="0.01 0.1 0.01"/>
|
39
|
-
<omc pos="0.0 0.0 -0.02" name="
|
40
|
-
<body joint="frozen" name="
|
39
|
+
<omc pos="0.0 0.0 -0.02" name="seg3" pos_marker="1"/>
|
40
|
+
<body joint="frozen" name="imu3_3Seg" pos="0.099999994 0.0 0.035" pos_min="0.050000012 -0.05 -0.05" pos_max="0.15 0.05 0.05">
|
41
41
|
<geom mass="0.1" color="dustin_exp_orange" edge_color="black" type="box" dim="0.05 0.03 0.02"/>
|
42
|
-
<omc pos="0.1 0.0 0.015" name="
|
42
|
+
<omc pos="0.1 0.0 0.015" name="seg3" pos_marker="1"/>
|
43
43
|
</body>
|
44
|
-
<body joint="ry" name="
|
44
|
+
<body joint="ry" name="seg4_3Seg" pos="0.20000002 0.0 0.0" pos_min="0.0 -0.05 -0.05" pos_max="0.35 0.05 0.05" damping="3.0">
|
45
45
|
<geom pos="0.1 0.0 0.0" mass="1.0" color="dustin_exp_white" edge_color="black" type="box" dim="0.2 0.05 0.05"/>
|
46
46
|
<geom pos="0.1 0.05 0.0" mass="0.1" color="black" edge_color="black" type="box" dim="0.01 0.1 0.01"/>
|
47
47
|
<geom pos="0.15 -0.05 0.0" mass="0.1" color="black" edge_color="black" type="box" dim="0.01 0.1 0.01"/>
|
48
|
-
<omc pos="0.0 0.0 -0.02" name="
|
49
|
-
<body joint="frozen" name="
|
48
|
+
<omc pos="0.0 0.0 -0.02" name="seg4" pos_marker="2"/>
|
49
|
+
<body joint="frozen" name="imu4_3Seg" pos="0.100000024 0.0 0.035" pos_min="0.050000012 -0.05 -0.05" pos_max="0.14999998 0.05 0.05">
|
50
50
|
<geom mass="0.1" color="dustin_exp_orange" edge_color="black" type="box" dim="0.05 0.03 0.02"/>
|
51
|
-
<omc pos="0.1 0.0 0.015" name="
|
51
|
+
<omc pos="0.1 0.0 0.015" name="seg4" pos_marker="2"/>
|
52
52
|
</body>
|
53
|
-
<body joint="rz" name="
|
53
|
+
<body joint="rz" name="seg5_3Seg" pos="0.19999999 0.0 0.0" pos_min="0.0 -0.05 -0.05" pos_max="0.35000002 0.05 0.05" damping="3.0">
|
54
54
|
<geom pos="0.1 0.0 0.0" mass="1.0" color="dustin_exp_white" edge_color="black" type="box" dim="0.2 0.05 0.05"/>
|
55
55
|
<geom pos="0.03 -0.05 0.0" mass="0.1" color="black" edge_color="black" type="box" dim="0.01 0.1 0.01"/>
|
56
56
|
<geom pos="0.17 -0.05 0.0" mass="0.1" color="black" edge_color="black" type="box" dim="0.01 0.1 0.01"/>
|
57
|
-
<omc pos="0.0 0.0 -0.02" name="
|
58
|
-
<body joint="frozen" name="
|
57
|
+
<omc pos="0.0 0.0 -0.02" name="seg5" pos_marker="4"/>
|
58
|
+
<body joint="frozen" name="imu5_3Seg" pos="0.100000024 0.0 0.035" pos_min="0.050000012 -0.05 -0.05" pos_max="0.15000004 0.05 0.05">
|
59
59
|
<geom mass="0.1" color="dustin_exp_orange" edge_color="black" type="box" dim="0.05 0.03 0.02"/>
|
60
|
-
<omc pos="0.1 0.0 0.015" name="
|
60
|
+
<omc pos="0.1 0.0 0.015" name="seg5" pos_marker="4"/>
|
61
61
|
</body>
|
62
62
|
</body>
|
63
63
|
</body>
|
64
64
|
</body>
|
65
|
-
<body joint="free" name="
|
65
|
+
<body joint="free" name="seg2_4Seg" pos="0.2 0.0 0.0" pos_min="0.15 -0.05 -0.05" pos_max="0.35 0.05 0.05" damping="5.0 5.0 5.0 25.0 25.0 25.0">
|
66
66
|
<geom pos="0.1 0.0 0.0" mass="1.0" color="dustin_exp_white" edge_color="black" type="box" dim="0.2 0.05 0.05"/>
|
67
67
|
<geom pos="0.03 -0.05 0.0" mass="0.1" color="dustin_exp_white" edge_color="black" type="box" dim="0.01 0.1 0.01"/>
|
68
68
|
<geom pos="0.17 -0.05 0.0" mass="0.1" color="dustin_exp_white" edge_color="black" type="box" dim="0.01 0.1 0.01"/>
|
69
|
-
<omc pos="0.0 0.0 -0.02" name="
|
70
|
-
<body joint="frozen" name="
|
69
|
+
<omc pos="0.0 0.0 -0.02" name="seg2" pos_marker="2"/>
|
70
|
+
<body joint="frozen" name="imu2_4Seg" pos="0.10000001 0.0 0.035" pos_min="0.049999997 -0.05 -0.05" pos_max="0.15000002 0.05 0.05">
|
71
71
|
<geom mass="0.1" color="dustin_exp_orange" edge_color="black" type="box" dim="0.05 0.03 0.02"/>
|
72
|
-
<omc pos="0.1 0.0 0.015" name="
|
72
|
+
<omc pos="0.1 0.0 0.015" name="seg2" pos_marker="2"/>
|
73
73
|
</body>
|
74
|
-
<body joint="rx" name="
|
74
|
+
<body joint="rx" name="seg3_4Seg" pos="0.2 0.0 0.0" pos_min="0.0 -0.05 -0.05" pos_max="0.35000002 0.05 0.05" damping="3.0">
|
75
75
|
<geom pos="0.1 0.0 0.0" mass="1.0" color="dustin_exp_blue" edge_color="black" type="box" dim="0.2 0.05 0.05"/>
|
76
76
|
<geom pos="0.05 0.05 0.0" mass="0.1" color="black" edge_color="black" type="box" dim="0.01 0.1 0.01"/>
|
77
77
|
<geom pos="0.15 -0.05 0.0" mass="0.1" color="black" edge_color="black" type="box" dim="0.01 0.1 0.01"/>
|
78
|
-
<omc pos="0.0 0.0 -0.02" name="
|
79
|
-
<body joint="frozen" name="
|
78
|
+
<omc pos="0.0 0.0 -0.02" name="seg3" pos_marker="1"/>
|
79
|
+
<body joint="frozen" name="imu3_4Seg" pos="0.099999994 0.0 0.035" pos_min="0.050000012 -0.05 -0.05" pos_max="0.15 0.05 0.05">
|
80
80
|
<geom mass="0.1" color="dustin_exp_orange" edge_color="black" type="box" dim="0.05 0.03 0.02"/>
|
81
|
-
<omc pos="0.1 0.0 0.015" name="
|
81
|
+
<omc pos="0.1 0.0 0.015" name="seg3" pos_marker="1"/>
|
82
82
|
</body>
|
83
|
-
<body joint="ry" name="
|
83
|
+
<body joint="ry" name="seg4_4Seg" pos="0.20000002 0.0 0.0" pos_min="0.0 -0.05 -0.05" pos_max="0.35 0.05 0.05" damping="3.0">
|
84
84
|
<geom pos="0.1 0.0 0.0" mass="1.0" color="dustin_exp_white" edge_color="black" type="box" dim="0.2 0.05 0.05"/>
|
85
85
|
<geom pos="0.1 0.05 0.0" mass="0.1" color="black" edge_color="black" type="box" dim="0.01 0.1 0.01"/>
|
86
86
|
<geom pos="0.15 -0.05 0.0" mass="0.1" color="black" edge_color="black" type="box" dim="0.01 0.1 0.01"/>
|
87
|
-
<omc pos="0.0 0.0 -0.02" name="
|
88
|
-
<body joint="frozen" name="
|
87
|
+
<omc pos="0.0 0.0 -0.02" name="seg4" pos_marker="2"/>
|
88
|
+
<body joint="frozen" name="imu4_4Seg" pos="0.100000024 0.0 0.035" pos_min="0.050000012 -0.05 -0.05" pos_max="0.14999998 0.05 0.05">
|
89
89
|
<geom mass="0.1" color="dustin_exp_orange" edge_color="black" type="box" dim="0.05 0.03 0.02"/>
|
90
|
-
<omc pos="0.1 0.0 0.015" name="
|
90
|
+
<omc pos="0.1 0.0 0.015" name="seg4" pos_marker="2"/>
|
91
91
|
</body>
|
92
|
-
<body joint="rz" name="
|
92
|
+
<body joint="rz" name="seg5_4Seg" pos="0.19999999 0.0 0.0" pos_min="0.0 -0.05 -0.05" pos_max="0.35000002 0.05 0.05" damping="3.0">
|
93
93
|
<geom pos="0.1 0.0 0.0" mass="1.0" color="dustin_exp_white" edge_color="black" type="box" dim="0.2 0.05 0.05"/>
|
94
94
|
<geom pos="0.03 -0.05 0.0" mass="0.1" color="black" edge_color="black" type="box" dim="0.01 0.1 0.01"/>
|
95
95
|
<geom pos="0.17 -0.05 0.0" mass="0.1" color="black" edge_color="black" type="box" dim="0.01 0.1 0.01"/>
|
96
|
-
<omc pos="0.0 0.0 -0.02" name="
|
97
|
-
<body joint="frozen" name="
|
96
|
+
<omc pos="0.0 0.0 -0.02" name="seg5" pos_marker="4"/>
|
97
|
+
<body joint="frozen" name="imu5_4Seg" pos="0.100000024 0.0 0.035" pos_min="0.050000012 -0.05 -0.05" pos_max="0.15000004 0.05 0.05">
|
98
98
|
<geom mass="0.1" color="dustin_exp_orange" edge_color="black" type="box" dim="0.05 0.03 0.02"/>
|
99
|
-
<omc pos="0.1 0.0 0.015" name="
|
99
|
+
<omc pos="0.1 0.0 0.015" name="seg5" pos_marker="4"/>
|
100
100
|
</body>
|
101
101
|
</body>
|
102
102
|
</body>
|