From 1c5293f341baa023ee78c94a3fc150b4ea29f96c Mon Sep 17 00:00:00 2001 From: HipsterBrown Date: Fri, 11 Sep 2026 16:34:37 -0400 Subject: [PATCH] Add EndPosition playback via a method config attribute New optional `method` attribute selects which arm capture to replay: JointPositions (default, MoveToJointPositions) or EndPosition (MoveToPosition). Pose payloads are decoded with protojson. The player now drives frames through a move callback so it is agnostic to the arm call. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01Gy9SjPAUdKsoA3VaMssg97 --- config_test.go | 10 ++++ devrel_sequence-playback_arm-playback.md | 6 +- go.mod | 4 +- module.go | 51 +++++++++++------ module_test.go | 27 +++++++++ playback.go | 15 +++-- playback_test.go | 47 ++++++++++++++-- sequence.go | 72 +++++++++++++++++------- sequence_test.go | 33 ++++++++++- 9 files changed, 210 insertions(+), 55 deletions(-) diff --git a/config_test.go b/config_test.go index 1acdba6..2513ae1 100644 --- a/config_test.go +++ b/config_test.go @@ -29,6 +29,16 @@ func TestConfigValidate(t *testing.T) { }) } +func TestConfigMethod(t *testing.T) { + for _, m := range []string{"", "JointPositions", "EndPosition"} { + _, _, err := (&Config{Arm: "a", Method: m}).Validate("x") + test.That(t, err, test.ShouldBeNil) + } + _, _, err := (&Config{Arm: "a", Method: "Pose"}).Validate("x") + test.That(t, err, test.ShouldNotBeNil) + test.That(t, err.Error(), test.ShouldContainSubstring, "method") +} + func TestConfigArmExtraDefault(t *testing.T) { test.That(t, (&Config{}).armExtra(), test.ShouldResemble, map[string]any{"wait": false, "streamed": true}) test.That(t, (&Config{ArmExtra: map[string]any{}}).armExtra(), test.ShouldResemble, map[string]any{}) diff --git a/devrel_sequence-playback_arm-playback.md b/devrel_sequence-playback_arm-playback.md index b8e6318..e39f826 100644 --- a/devrel_sequence-playback_arm-playback.md +++ b/devrel_sequence-playback_arm-playback.md @@ -17,6 +17,7 @@ Replays a recorded sequence on an arm, and optionally a gripper, at the recorded | Name | Type | Required | Default | Description | |---|---|---|---|---| | `arm` | string | Yes | | The arm to move. | +| `method` | string | No | `JointPositions` | Which arm capture to replay: `JointPositions` (replayed with `MoveToJointPositions`) or `EndPosition` (replayed with `MoveToPosition`). | | `gripper` | string | No | | The gripper to move, using its `set_position` DoCommand. | | `gripper_position_key` | string | No | `position_percentage` | The field in the gripper's `get_position` response that holds the position. Also used when calling `set_position`. | | `source_arm` | string | No | same as `arm` | The name of the arm in the recording. Set this to the leader arm to replay its motion on the follower. | @@ -26,7 +27,8 @@ Replays a recorded sequence on an arm, and optionally a gripper, at the recorded ## What the recording needs to contain -- For the arm: `JointPositions` captured by the data manager. +- For the arm: `JointPositions` (default) or `EndPosition` captured by the data manager, + matching `method`. - For the gripper: a `DoCommand` capture with `{"command": "get_position"}` as the input. ## How playback works @@ -55,7 +57,7 @@ Returns `sequences`, each with `id`, `tags`, `start`, and `duration_s`. ### play Starts playing a sequence in the background and returns right away. Fails if something is -already playing or if the recording's joint count doesn't match the arm. +already playing or, for `JointPositions`, if the recording's joint count doesn't match the arm. ```json { "command": "play", "sequence_id": "" } diff --git a/go.mod b/go.mod index 1b01a20..1008c40 100644 --- a/go.mod +++ b/go.mod @@ -3,10 +3,12 @@ module sequenceplayback go 1.25.10 require ( + github.com/golang/geo v0.0.0-20230421003525-6adc56603217 go.viam.com/api v0.1.580 go.viam.com/rdk v1.7.0 go.viam.com/test v1.2.4 golang.org/x/sync v0.21.0 + google.golang.org/protobuf v1.36.11 ) require ( @@ -90,7 +92,6 @@ require ( github.com/gogo/protobuf v1.3.2 // indirect github.com/golang-jwt/jwt/v4 v4.5.2 // indirect github.com/golang/freetype v0.0.0-20170609003504-e2365dfdc4a0 // indirect - github.com/golang/geo v0.0.0-20230421003525-6adc56603217 // indirect github.com/golang/groupcache v0.0.0-20241129210726-2c02b8208cf8 // indirect github.com/golang/protobuf v1.5.4 // indirect github.com/golang/snappy v0.0.4 // indirect @@ -220,7 +221,6 @@ require ( google.golang.org/genproto/googleapis/api v0.0.0-20260610212136-7ab31c22f7ad // indirect google.golang.org/genproto/googleapis/rpc v0.0.0-20260610212136-7ab31c22f7ad // indirect google.golang.org/grpc v1.83.1 // indirect - google.golang.org/protobuf v1.36.11 // indirect gopkg.in/natefinch/lumberjack.v2 v2.2.1 // indirect gopkg.in/yaml.v3 v3.0.1 // indirect gorgonia.org/tensor v0.9.24 // indirect diff --git a/module.go b/module.go index d2ad1b2..638cb5f 100644 --- a/module.go +++ b/module.go @@ -29,6 +29,7 @@ func init() { type Config struct { Arm string `json:"arm"` + Method string `json:"method,omitempty"` // JointPositions (default) or EndPosition Gripper string `json:"gripper,omitempty"` GripperPositionKey string `json:"gripper_position_key,omitempty"` SourceArm string `json:"source_arm,omitempty"` @@ -44,6 +45,11 @@ func (cfg *Config) Validate(path string) ([]string, []string, error) { if cfg.SourceGripper != "" && cfg.Gripper == "" { return nil, nil, fmt.Errorf("%s: source_gripper requires gripper", path) } + switch cfg.Method { + case "", "JointPositions", "EndPosition": + default: + return nil, nil, fmt.Errorf("%s: method must be JointPositions or EndPosition, got %q", path, cfg.Method) + } deps := []string{cfg.Arm} if cfg.Gripper != "" { deps = append(deps, cfg.Gripper) @@ -118,7 +124,8 @@ func NewArmPlayback(ctx context.Context, deps resource.Dependencies, rawConf res if s.cloudErr != nil { return nil, s.cloudErr } - return fetchSequence(ctx, s.dc, id, cmp.Or(cfg.SourceArm, cfg.Arm), cmp.Or(cfg.SourceGripper, cfg.Gripper), s.player.gripperKey) + return fetchSequence(ctx, s.dc, id, cmp.Or(cfg.Method, "JointPositions"), + cmp.Or(cfg.SourceArm, cfg.Arm), cmp.Or(cfg.SourceGripper, cfg.Gripper), s.player.gripperKey) } return s, nil } @@ -183,16 +190,26 @@ func (s *armPlayback) play(ctx context.Context, cmd map[string]any) (map[string] return nil, err } } - current, err := s.arm.JointPositions(ctx, nil) - if err != nil { - return nil, fmt.Errorf("reading arm joints: %w", err) - } - if len(current) != len(seq.JointsDeg[0]) { - return nil, fmt.Errorf("joint count mismatch: arm has %d, sequence %q has %d", len(current), id, len(seq.JointsDeg[0])) - } - inputs, err := s.toInputs(ctx, seq.JointsDeg) - if err != nil { - return nil, err + var move moveFunc + if seq.HasPoses() { + move = func(ctx context.Context, i int, extra map[string]any) error { + return s.arm.MoveToPosition(ctx, seq.Poses[i], extra) + } + } else { + current, err := s.arm.JointPositions(ctx, nil) + if err != nil { + return nil, fmt.Errorf("reading arm joints: %w", err) + } + if len(current) != len(seq.JointsDeg[0]) { + return nil, fmt.Errorf("joint count mismatch: arm has %d, sequence %q has %d", len(current), id, len(seq.JointsDeg[0])) + } + inputs, err := s.toInputs(ctx, seq.JointsDeg) + if err != nil { + return nil, err + } + move = func(ctx context.Context, i int, extra map[string]any) error { + return s.arm.MoveToJointPositions(ctx, inputs[i], extra) + } } useGripper := s.player.gripper != nil && seq.HasGripper() if s.player.gripper != nil && !seq.HasGripper() { @@ -209,16 +226,16 @@ func (s *armPlayback) play(ctx context.Context, cmd map[string]any) (map[string] s.lastError = "" wctx, cancel := context.WithCancel(context.Background()) s.cancel, s.done = cancel, make(chan struct{}) - go s.worker(wctx, seq, inputs) - s.logger.Infow("playback started", "sequence_id", id, "frames", len(inputs), "duration", seq.Duration()) + go s.worker(wctx, seq, move) + s.logger.Infow("playback started", "sequence_id", id, "frames", seq.Len(), "duration", seq.Duration()) return map[string]any{ - "status": "playing", "sequence_id": id, "frame_count": len(inputs), + "status": "playing", "sequence_id": id, "frame_count": seq.Len(), "duration_s": seq.Duration().Seconds(), "has_gripper": useGripper, }, nil } -func (s *armPlayback) worker(ctx context.Context, seq *Sequence, inputs [][]referenceframe.Input) { - err := s.player.run(ctx, seq, inputs) +func (s *armPlayback) worker(ctx context.Context, seq *Sequence, move moveFunc) { + err := s.player.run(ctx, seq, move) s.mu.Lock() defer s.mu.Unlock() switch { @@ -252,7 +269,7 @@ func (s *armPlayback) Status(context.Context) (map[string]any, error) { st := map[string]any{"state": "idle", "sequence_id": s.sequence, "last_error": s.lastError, "frame_index": int(s.player.frameIndex.Load())} if seq := s.cache[s.sequence]; seq != nil { - st["frame_count"] = len(seq.JointsDeg) + st["frame_count"] = seq.Len() } if s.cancel != nil { st["state"] = "playing" diff --git a/module_test.go b/module_test.go index 8d59216..ab430e6 100644 --- a/module_test.go +++ b/module_test.go @@ -13,6 +13,7 @@ import ( "go.viam.com/rdk/referenceframe" "go.viam.com/rdk/resource" generic "go.viam.com/rdk/services/generic" + "go.viam.com/rdk/spatialmath" "go.viam.com/rdk/testutils/inject" "go.viam.com/test" ) @@ -170,6 +171,32 @@ func TestPlayErrorSurfacesThenClears(t *testing.T) { test.That(t, st["last_error"], test.ShouldBeEmpty) } +func TestPlayEndPositionUsesMoveToPosition(t *testing.T) { + r := newSvc(t, &Config{Arm: "a", Method: "EndPosition"}) + // No joints on the arm at all: the pose path must not consult them. + r.arm.JointPositionsFunc = func(context.Context, map[string]any) ([]referenceframe.Input, error) { + return nil, errors.New("not called") + } + var got []spatialmath.Pose + r.arm.MoveToPositionFunc = func(_ context.Context, p spatialmath.Pose, _ map[string]any) error { + got = append(got, p) + return nil + } + r.svc.fetch = func(_ context.Context, id string) (*Sequence, error) { + s := threeFrames() + s.JointsDeg = nil + s.Poses = []spatialmath.Pose{spatialmath.NewZeroPose(), spatialmath.NewZeroPose(), spatialmath.NewZeroPose()} + return s, nil + } + resp, err := do(t, r.svc, map[string]any{"command": "play", "sequence_id": "p1"}) + test.That(t, err, test.ShouldBeNil) + test.That(t, resp["frame_count"], test.ShouldEqual, 3) + st := waitIdle(t, r.svc) + test.That(t, st["last_error"], test.ShouldBeEmpty) + test.That(t, st["frame_count"], test.ShouldEqual, 3) + test.That(t, len(got), test.ShouldEqual, 3) +} + func TestListRequiresDatasetID(t *testing.T) { r := newSvc(t, &Config{Arm: "a"}) _, err := do(t, r.svc, map[string]any{"command": "list"}) diff --git a/playback.go b/playback.go index 10f2416..0554a8f 100644 --- a/playback.go +++ b/playback.go @@ -8,7 +8,6 @@ import ( "go.viam.com/rdk/components/arm" "go.viam.com/rdk/components/gripper" - "go.viam.com/rdk/referenceframe" "golang.org/x/sync/errgroup" ) @@ -52,9 +51,13 @@ func (p *player) setGripper(ctx context.Context, pos float64) error { return err } +// moveFunc issues arm frame i with the given extra; the caller decides whether that is a +// joint or a pose move. +type moveFunc func(ctx context.Context, i int, extra map[string]any) error + // run blocks until the sequence has played, failed, or ctx is cancelled. The arm is always -// sent Stop before returning. inputs are seq.JointsDeg already converted to arm inputs. -func (p *player) run(ctx context.Context, seq *Sequence, inputs [][]referenceframe.Input) (err error) { +// sent Stop before returning. +func (p *player) run(ctx context.Context, seq *Sequence, move moveFunc) (err error) { p.frameIndex.Store(0) defer func() { if stopErr := p.arm.Stop(context.Background(), nil); stopErr != nil && err == nil { @@ -66,7 +69,7 @@ func (p *player) run(ctx context.Context, seq *Sequence, inputs [][]referencefra // Safe entry: a blocking move to the first frame, and the gripper to its first value. g, gctx := errgroup.WithContext(ctx) - g.Go(func() error { return p.arm.MoveToJointPositions(gctx, inputs[0], nil) }) + g.Go(func() error { return move(gctx, 0, nil) }) if useGripper { g.Go(func() error { return p.setGripper(gctx, seq.GripperPositions[0]) }) } @@ -82,11 +85,11 @@ func (p *player) run(ctx context.Context, seq *Sequence, inputs [][]referencefra } g, gctx = errgroup.WithContext(ctx) g.Go(func() error { - for i := 1; i < len(inputs); i++ { + for i := 1; i < seq.Len(); i++ { if err := waitUntil(gctx, seq.Offsets[i]); err != nil { return err } - if err := p.arm.MoveToJointPositions(gctx, inputs[i], p.armExtra); err != nil { + if err := move(gctx, i, p.armExtra); err != nil { return fmt.Errorf("frame %d: %w", i, err) } p.frameIndex.Store(int64(i + 1)) diff --git a/playback_test.go b/playback_test.go index d2e84b1..052c6ee 100644 --- a/playback_test.go +++ b/playback_test.go @@ -2,12 +2,15 @@ package sequenceplayback import ( "context" + "errors" + "github.com/golang/geo/r3" "sync" "testing" "time" "go.viam.com/rdk/referenceframe" + "go.viam.com/rdk/spatialmath" "go.viam.com/rdk/testutils/inject" "go.viam.com/test" ) @@ -89,6 +92,13 @@ func threeFrames() *Sequence { } } +// joints adapts converted inputs to the player's per-frame move callback. +func (r *rig) joints(inputs [][]referenceframe.Input) moveFunc { + return func(ctx context.Context, i int, extra map[string]any) error { + return r.arm.MoveToJointPositions(ctx, inputs[i], extra) + } +} + func rads(deg ...float64) []referenceframe.Input { out := make([]referenceframe.Input, len(deg)) for i, d := range deg { @@ -100,7 +110,7 @@ func rads(deg ...float64) []referenceframe.Input { func TestRunPacesArmFramesByRecordedOffsets(t *testing.T) { r := newRig(t, false) inputs := [][]referenceframe.Input{rads(0, 0), rads(90, 0), rads(180, 0)} - err := r.player.run(context.Background(), threeFrames(), inputs) + err := r.player.run(context.Background(), threeFrames(), r.joints(inputs)) test.That(t, err, test.ShouldBeNil) test.That(t, len(r.moves), test.ShouldEqual, 3) @@ -122,7 +132,7 @@ func TestRunDrivesGripperOnSharedClockAndDedupes(t *testing.T) { seq.GripperOffsets = []time.Duration{-20 * time.Millisecond, 100 * time.Millisecond, 200 * time.Millisecond, 250 * time.Millisecond} seq.GripperPositions = []float64{90, 90, 10, 10} inputs := [][]referenceframe.Input{rads(0, 0), rads(90, 0), rads(180, 0)} - err := r.player.run(context.Background(), seq, inputs) + err := r.player.run(context.Background(), seq, r.joints(inputs)) test.That(t, err, test.ShouldBeNil) // Safe entry sets the first value; the repeated 90 and repeated 10 are skipped. @@ -136,7 +146,7 @@ func TestRunDrivesGripperOnSharedClockAndDedupes(t *testing.T) { func TestRunWithoutGripperDataLeavesGripperAlone(t *testing.T) { r := newRig(t, true) inputs := [][]referenceframe.Input{rads(0, 0), rads(90, 0), rads(180, 0)} - err := r.player.run(context.Background(), threeFrames(), inputs) + err := r.player.run(context.Background(), threeFrames(), r.joints(inputs)) test.That(t, err, test.ShouldBeNil) test.That(t, r.grips, test.ShouldBeEmpty) } @@ -153,7 +163,7 @@ func TestRunStopsArmOnMoveError(t *testing.T) { return nil } inputs := [][]referenceframe.Input{rads(0, 0), rads(90, 0), rads(180, 0)} - err := r.player.run(context.Background(), threeFrames(), inputs) + err := r.player.run(context.Background(), threeFrames(), r.joints(inputs)) test.That(t, errors.Is(err, boom), test.ShouldBeTrue) test.That(t, n, test.ShouldEqual, 2) test.That(t, r.stops, test.ShouldEqual, 1) @@ -167,7 +177,7 @@ func TestRunStopsArmOnCancel(t *testing.T) { return nil } inputs := [][]referenceframe.Input{rads(0, 0), rads(90, 0), rads(180, 0)} - err := r.player.run(ctx, threeFrames(), inputs) + err := r.player.run(ctx, threeFrames(), r.joints(inputs)) test.That(t, errors.Is(err, context.Canceled), test.ShouldBeTrue) test.That(t, r.stops, test.ShouldEqual, 1) } @@ -181,9 +191,34 @@ func TestRunSafeEntryGripperErrorAborts(t *testing.T) { seq.GripperOffsets = []time.Duration{0} seq.GripperPositions = []float64{50} inputs := [][]referenceframe.Input{rads(0, 0), rads(90, 0), rads(180, 0)} - err := r.player.run(context.Background(), seq, inputs) + err := r.player.run(context.Background(), seq, r.joints(inputs)) test.That(t, err, test.ShouldNotBeNil) test.That(t, err.Error(), test.ShouldContainSubstring, "gripper offline") test.That(t, len(r.moves), test.ShouldBeLessThanOrEqualTo, 1) test.That(t, r.stops, test.ShouldEqual, 1) } + +func TestRunPacesPoseFrames(t *testing.T) { + r := newRig(t, false) + var poses []spatialmath.Pose + var at []time.Time + r.arm.MoveToPositionFunc = func(_ context.Context, p spatialmath.Pose, _ map[string]any) error { + poses = append(poses, p) + at = append(at, r.clock.Now()) + return nil + } + seq := threeFrames() + seq.JointsDeg = nil + for i := range 3 { + seq.Poses = append(seq.Poses, spatialmath.NewPoseFromPoint(r3.Vector{X: float64(i)})) + } + move := func(ctx context.Context, i int, extra map[string]any) error { + return r.arm.MoveToPosition(ctx, seq.Poses[i], extra) + } + err := r.player.run(context.Background(), seq, move) + test.That(t, err, test.ShouldBeNil) + test.That(t, poses, test.ShouldResemble, seq.Poses) + test.That(t, r.moves, test.ShouldBeEmpty) + test.That(t, at[2].Sub(at[0]), test.ShouldEqual, 300*time.Millisecond) + test.That(t, r.stops, test.ShouldEqual, 1) +} diff --git a/sequence.go b/sequence.go index 0e74f23..6e5f0f5 100644 --- a/sequence.go +++ b/sequence.go @@ -2,26 +2,34 @@ package sequenceplayback import ( "context" + "encoding/json" "fmt" "sort" "time" + commonpb "go.viam.com/api/common/v1" "go.viam.com/rdk/app" "go.viam.com/rdk/components/arm" "go.viam.com/rdk/components/gripper" + "go.viam.com/rdk/spatialmath" + "google.golang.org/protobuf/encoding/protojson" ) -// Sequence is one captured demonstration, ready to replay. Joint values are kept in the +// Sequence is one captured demonstration, ready to replay. Exactly one of JointsDeg or +// Poses is populated, depending on the captured arm method. Joint values are kept in the // captured unit (degrees); conversion to the arm's inputs happens at play time. type Sequence struct { ID string - Offsets []time.Duration // per frame, from the first arm row - JointsDeg [][]float64 - GripperOffsets []time.Duration // same origin as Offsets; may be negative + Offsets []time.Duration // per frame, from the first arm row + JointsDeg [][]float64 // JointPositions captures + Poses []spatialmath.Pose // EndPosition captures + GripperOffsets []time.Duration // same origin as Offsets; may be negative GripperPositions []float64 } func (s *Sequence) HasGripper() bool { return len(s.GripperPositions) > 0 } +func (s *Sequence) HasPoses() bool { return len(s.Poses) > 0 } +func (s *Sequence) Len() int { return len(s.Offsets) } func (s *Sequence) Duration() time.Duration { return s.Offsets[len(s.Offsets)-1] } func sortRows(rows []*app.ExportTabularDataResponse) { @@ -45,6 +53,24 @@ func jointValues(payload map[string]any) ([]float64, error) { return out, nil } +// poseValue decodes an EndPosition payload. protojson accepts both the exported snake_case +// keys (o_x) and lowerCamel (oX). +func poseValue(payload map[string]any) (spatialmath.Pose, error) { + raw, ok := payload["pose"] + if !ok { + return nil, fmt.Errorf("payload has no pose") + } + b, err := json.Marshal(raw) + if err != nil { + return nil, err + } + var p commonpb.Pose + if err := (protojson.UnmarshalOptions{DiscardUnknown: true}).Unmarshal(b, &p); err != nil { + return nil, fmt.Errorf("decoding pose: %w", err) + } + return spatialmath.NewPoseFromProtobuf(&p), nil +} + func gripperValue(payload map[string]any, key string) (float64, error) { out, _ := payload["docommand_output"].(map[string]any) v, ok := out[key] @@ -58,11 +84,11 @@ func gripperValue(payload map[string]any, key string) (float64, error) { return f, nil } -// buildSequence turns exported JointPositions rows (and optional gripper DoCommand rows) -// into a Sequence. Rows may arrive in any order. -func buildSequence(id string, armRows, gripperRows []*app.ExportTabularDataResponse, gripperKey string) (*Sequence, error) { +// buildSequence turns exported arm rows (JointPositions or EndPosition, per method) and +// optional gripper DoCommand rows into a Sequence. Rows may arrive in any order. +func buildSequence(id, method string, armRows, gripperRows []*app.ExportTabularDataResponse, gripperKey string) (*Sequence, error) { if len(armRows) == 0 { - return nil, fmt.Errorf("sequence %q has no JointPositions rows", id) + return nil, fmt.Errorf("sequence %q has no %s rows", id, method) } sortRows(armRows) sortRows(gripperRows) @@ -70,14 +96,22 @@ func buildSequence(id string, armRows, gripperRows []*app.ExportTabularDataRespo seq := &Sequence{ID: id} for i, r := range armRows { - j, err := jointValues(r.Payload) - if err != nil { - return nil, fmt.Errorf("sequence %q arm row %d: %w", id, i, err) - } - if i > 0 && len(j) != len(seq.JointsDeg[0]) { - return nil, fmt.Errorf("sequence %q arm row %d has %d joints, expected %d (joint count changed mid-sequence)", id, i, len(j), len(seq.JointsDeg[0])) + if method == "EndPosition" { + p, err := poseValue(r.Payload) + if err != nil { + return nil, fmt.Errorf("sequence %q arm row %d: %w", id, i, err) + } + seq.Poses = append(seq.Poses, p) + } else { + j, err := jointValues(r.Payload) + if err != nil { + return nil, fmt.Errorf("sequence %q arm row %d: %w", id, i, err) + } + if i > 0 && len(j) != len(seq.JointsDeg[0]) { + return nil, fmt.Errorf("sequence %q arm row %d has %d joints, expected %d (joint count changed mid-sequence)", id, i, len(j), len(seq.JointsDeg[0])) + } + seq.JointsDeg = append(seq.JointsDeg, j) } - seq.JointsDeg = append(seq.JointsDeg, j) seq.Offsets = append(seq.Offsets, r.TimeCaptured.Sub(origin)) } for i, r := range gripperRows { @@ -92,15 +126,15 @@ func buildSequence(id string, armRows, gripperRows []*app.ExportTabularDataRespo } // fetchSequence pulls one captured sequence from Viam Cloud. -func fetchSequence(ctx context.Context, dc *app.DataClient, id, sourceArm, sourceGripper, gripperKey string) (*Sequence, error) { +func fetchSequence(ctx context.Context, dc *app.DataClient, id, method, sourceArm, sourceGripper, gripperKey string) (*Sequence, error) { meta, err := dc.GetSequence(ctx, id) if err != nil { return nil, fmt.Errorf("getting sequence %q: %w", id, err) } interval := app.CaptureInterval{Start: meta.StartTime, End: meta.EndTime} - armRows, err := dc.ExportTabularData(ctx, meta.PartID, sourceArm, arm.API.String(), "JointPositions", interval, nil) + armRows, err := dc.ExportTabularData(ctx, meta.PartID, sourceArm, arm.API.String(), method, interval, nil) if err != nil { - return nil, fmt.Errorf("exporting %s JointPositions for sequence %q: %w", sourceArm, id, err) + return nil, fmt.Errorf("exporting %s %s for sequence %q: %w", sourceArm, method, id, err) } var gripperRows []*app.ExportTabularDataResponse if sourceGripper != "" { @@ -109,7 +143,7 @@ func fetchSequence(ctx context.Context, dc *app.DataClient, id, sourceArm, sourc return nil, fmt.Errorf("exporting %s DoCommand for sequence %q: %w", sourceGripper, id, err) } } - return buildSequence(id, armRows, gripperRows, gripperKey) + return buildSequence(id, method, armRows, gripperRows, gripperKey) } // listSequences summarizes every sequence in a dataset. diff --git a/sequence_test.go b/sequence_test.go index 8a199a7..68f313d 100644 --- a/sequence_test.go +++ b/sequence_test.go @@ -4,7 +4,9 @@ import ( "testing" "time" + "github.com/golang/geo/r3" "go.viam.com/rdk/app" + "go.viam.com/rdk/spatialmath" "go.viam.com/test" ) @@ -21,6 +23,13 @@ func armRow(at time.Duration, deg ...float64) *app.ExportTabularDataResponse { } } +func poseRow(at time.Duration, x float64) *app.ExportTabularDataResponse { + return &app.ExportTabularDataResponse{ + TimeCaptured: t0.Add(at), + Payload: map[string]any{"pose": map[string]any{"x": x, "y": 2.0, "z": 3.0, "o_x": 0.0, "o_y": 0.0, "o_z": 1.0, "theta": 90.0}}, + } +} + func gripRow(at time.Duration, key string, v any) *app.ExportTabularDataResponse { return &app.ExportTabularDataResponse{ TimeCaptured: t0.Add(at), @@ -38,7 +47,7 @@ func TestBuildSequenceSortsAndOffsetsFromFirstArmRow(t *testing.T) { gripRow(150*time.Millisecond, "position_percentage", 50.0), gripRow(-50*time.Millisecond, "position_percentage", 90.0), } - seq, err := buildSequence("abc", arm, grip, "position_percentage") + seq, err := buildSequence("abc", "JointPositions", arm, grip, "position_percentage") test.That(t, err, test.ShouldBeNil) test.That(t, seq.ID, test.ShouldEqual, "abc") test.That(t, seq.JointsDeg, test.ShouldResemble, [][]float64{{1, 2}, {5, 6}, {10, 20}}) @@ -49,11 +58,29 @@ func TestBuildSequenceSortsAndOffsetsFromFirstArmRow(t *testing.T) { } func TestBuildSequenceNoGripperRows(t *testing.T) { - seq, err := buildSequence("abc", []*app.ExportTabularDataResponse{armRow(0, 1)}, nil, "position_percentage") + seq, err := buildSequence("abc", "JointPositions", []*app.ExportTabularDataResponse{armRow(0, 1)}, nil, "position_percentage") test.That(t, err, test.ShouldBeNil) test.That(t, seq.HasGripper(), test.ShouldBeFalse) } +func TestBuildSequenceEndPosition(t *testing.T) { + rows := []*app.ExportTabularDataResponse{poseRow(100*time.Millisecond, 10), poseRow(0, 1)} + seq, err := buildSequence("abc", "EndPosition", rows, nil, "") + test.That(t, err, test.ShouldBeNil) + test.That(t, seq.HasPoses(), test.ShouldBeTrue) + test.That(t, seq.JointsDeg, test.ShouldBeNil) + test.That(t, seq.Len(), test.ShouldEqual, 2) + test.That(t, seq.Offsets, test.ShouldResemble, []time.Duration{0, 100 * time.Millisecond}) + want := spatialmath.NewPose(r3.Vector{X: 1, Y: 2, Z: 3}, &spatialmath.OrientationVectorDegrees{OZ: 1, Theta: 90}) + test.That(t, spatialmath.PoseAlmostEqual(seq.Poses[0], want), test.ShouldBeTrue) + test.That(t, seq.Poses[1].Point().X, test.ShouldEqual, 10.0) + + _, err = buildSequence("abc", "EndPosition", []*app.ExportTabularDataResponse{armRow(0, 1)}, nil, "") + test.That(t, err.Error(), test.ShouldContainSubstring, "no pose") + _, err = buildSequence("abc", "EndPosition", nil, nil, "") + test.That(t, err.Error(), test.ShouldContainSubstring, "no EndPosition rows") +} + func TestBuildSequenceErrors(t *testing.T) { cases := map[string]struct { arm, grip []*app.ExportTabularDataResponse @@ -69,7 +96,7 @@ func TestBuildSequenceErrors(t *testing.T) { } for name, c := range cases { t.Run(name, func(t *testing.T) { - _, err := buildSequence("abc", c.arm, c.grip, "position_percentage") + _, err := buildSequence("abc", "JointPositions", c.arm, c.grip, "position_percentage") test.That(t, err, test.ShouldNotBeNil) test.That(t, err.Error(), test.ShouldContainSubstring, c.want) })